Python in O(lg n) with max heap and min heap


  • 0
    A

    Inspired from here: http://www.geeksforgeeks.org/median-of-stream-of-integers-running-integers/

    use max heap represent half left and use min heap represent half right.

    Then you compare the root of both and the size of both heap. You can get the median

    def __init__(self):
        self.max_heap = [] # for half left
        self.min_heap = [] # for half right
    
    def addNum(self, num):
        size_left = len(self.max_heap)
        size_right = len(self.min_heap)
        if size_left == 0 and size_right == 0:
            heappush(self.min_heap, num)
            return
        
        if self.max_heap:
            left_val = -self.max_heap[0]
        if self.min_heap:
            right_val = self.min_heap[0]
        
        if size_left == size_right:
            if num <= left_val:
                heappush(self.max_heap, -num)
            else:
                heappush(self.min_heap, num)
        elif size_left > size_right:
            if num >= left_val:
                heappush(self.min_heap, num) # then balance
            else:
                tmp = heappop(self.max_heap)
                heappush(self.min_heap, -tmp)
                heappush(self.max_heap, -num)
        else: # size_left < size_right
            if num < right_val:
                heappush(self.max_heap, -num) # then balance
            else:
                tmp = heappop(self.min_heap)
                heappush(self.max_heap, -tmp)
                heappush(self.min_heap, num)       
    
    def findMedian(self):
        size_left = len(self.max_heap)
        size_right = len(self.min_heap)
        
        if size_left == size_right:
            return (-self.max_heap[0] + self.min_heap[0])/2.0
        elif size_left > size_right:
            return float(-self.max_heap[0])
        else:
            return float(self.min_heap[0])

Log in to reply
 

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