LC 0295 [H] Find Median from Data Stream - ALawliet/algorithms GitHub Wiki

balance the min heap first

class MedianFinder:
    def __init__(self):
        self.smallMaxHeap = []  # the smaller half of the list, max heap (invert min-heap)
        self.largeMinHeap = []  # the larger half of the list, min heap

    def addNum(self, num):
        if len(self.smallMaxHeap) == len(self.largeMinHeap):
            heappush(self.largeMinHeap, -heappushpop(self.smallMaxHeap, -num))
        else:
            heappush(self.smallMaxHeap, -heappushpop(self.largeMinHeap, num))

    def findMedian(self):
        if len(self.smallMaxHeap) == len(self.largeMinHeap):
            return (self.largeMinHeap[0] - self.smallMaxHeap[0]) / 2.0
        else:
            return float(self.largeMinHeap[0])
        


# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()
class MedianFinder:
    def __init__(self):
        self.maxHeap = []  # the smaller half of the list, max heap (invert min-heap)
        self.minHeap = []  # the larger half of the list, min heap

    def addNum(self, num):
        if not self.maxHeap or num <= -self.maxHeap[0]:
          heappush(self.maxHeap, -num)
        else:
          heappush(self.minHeap, num)

        # either both the heaps will have equal number of elements or max-heap will have one
        # more element than the min-heap
        if len(self.maxHeap) - len(self.minHeap) >= 2:
          heappush(self.minHeap, -heappop(self.maxHeap))
        elif len(self.maxHeap) < len(self.minHeap):
          heappush(self.maxHeap, -heappop(self.minHeap))

    def findMedian(self):
        if len(self.maxHeap) == len(self.minHeap):
          # we have even number of elements, take the average of middle two elements
          return (-self.maxHeap[0] + self.minHeap[0]) / 2.0

        # because max-heap will have one more element than the min-heap
        return -self.maxHeap[0] / 1.0



# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()