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()