CODE
class MaxHeap:
"""pop returns min"""
def __init__(self):
self.data = []
heapq.heapify(self.data)
def push(self, val):
heapq.heappush(self.data, -1 * val)
def pop(self):
return -1 * heapq.heappop(self.data)
def get_max(self):
return -1 * self.data[0] if self.data else float('-inf')
def __len__(self):
return len(self.data)
class MinHeap:
"""pop returns min"""
def __init__(self):
self.data = []
heapq.heapify(self.data)
def push(self, val):
heapq.heappush(self.data, val)
def pop(self):
return heapq.heappop(self.data)
def get_min(self):
return self.data[0] if self.data else float('inf')
def __len__(self):
return len(self.data)
class MedianFinder:
def __init__(self):
self.left = MaxHeap()
self.right = MinHeap()
def balance_heaps(self):
if len(self.left) > len(self.right):
self.right.push(self.left.pop())
elif len(self.left) < len(self.right):
self.left.push(self.right.pop())
def addNum(self, num: int) -> None:
if num > self.left.get_max():
self.right.push(num)
else:
self.left.push(num)
self.balance_heaps()
def findMedian(self) -> float:
if len(self.left) > len(self.right):
return self.left.get_max()
elif len(self.left) < len(self.right):
return self.right.get_min()
else:
return (self.left.get_max() + self.right.get_min()) / 2
# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()
Last updated