Python solution with detailed explanation


  • 0
    G

    Solution

    1. Visualize a sorted list. The middle two elements (M1, M2) determine the median. Element 0 to M1 is max heap. Element M2 to N is min heap.
    2. How to peek top element in heap? It is element at index 0. http://stackoverflow.com/questions/1750991/peeking-in-a-heap-in-python
    3. There are two APIs addNum and findMedian
    4. addNum(num)
    • First number added to max heap.
    • If a number is less than the top number in max heap, add to max heap. Otherwise to min-heap
    • If the two heaps are not balanced (size off by two), move top of one to other
    1. findMedian()
    • If the number of elements is even, return the average to top elements of both heaps
    • If not, then one heap has more elements than another. Return the top element of the heap
      with more elements.
    from heapq import heappush, heappop
    
    class MedianFinder:
        def __init__(self):
            """
            Initialize your data structure here.
            """
            self.heap_max, self.heap_min = [], []
            return
    
        def addNum(self, num):
            """
            Adds a num into the data structure.
            :type num: int
            :rtype: void
            """
            if len(self.heap_max) == 0 or self.heap_max[0]*-1 >= num:
                heappush(self.heap_max, num*-1)
            else:
                heappush(self.heap_min, num) 
            N1, N2 = len(self.heap_max), len(self.heap_min)
            if N1-N2 == 2:
                x = heappop(self.heap_max)
                heappush(self.heap_min, x*-1)
            elif N2-N1 == 2:
                x = heappop(self.heap_min)
                heappush(self.heap_max, x*-1)            
    
        def findMedian(self):
            """
            Returns the median of current data stream
            :rtype: float
            """
            N1, N2 = len(self.heap_max), len(self.heap_min)
            N = N1 + N2
            if N % 2 == 0:
                return (self.heap_max[0]*-1 + self.heap_min[0])/2.0
            else:
                if N1 > N2:
                    return self.heap_max[0]*-1
                else:
                    return self.heap_min[0]
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.