Find Median from Data Stream
- 2 min read
- LC-Hard
- Binary-heap
3 years ago
Solution
- Have two heaps: elements in left heap (small heap) are always lesser than or equal to all elements in the right heap (large heap)
- Size of two heaps always differ by at most 1 element at any time
- The left heap will be a max heap and the right heap will be a min heap
- Max heap because we can get the right most in the smaller heap, and vice versa for the larger heap. Median would be very easy to calculate
- By default, push to small heap. Check if value constraint is maintained. If not, move to large heap.
- Now, check if size constraint is maintained, and move elements between heaps accordingly
| Time | Space | Explanation |
|---|---|---|
O(log n) | O(1) | O(log n) for moving elements across heaps, O(1) for findMedian() |
def __init__(self):
self.small_heap = []
self.large_heap = []
def addNum(self, num: int) -> None:
heappush(self.small_heap, -num)
if self.large_heap and -self.small_heap[0] > self.large_heap[0]:
heappush(self.large_heap, -heappop(self.small_heap))
if len(self.small_heap) - len(self.large_heap) > 1:
heappush(self.large_heap, -heappop(self.small_heap))
elif len(self.large_heap) - len(self.small_heap) > 1:
heappush(self.small_heap, -heappop(self.large_heap))
def findMedian(self) -> float:
if len(self.small_heap) == len(self.large_heap):
return (-self.small_heap[0] + self.large_heap[0]) / 2
elif len(self.small_heap) > len(self.large_heap):
return -self.small_heap[0]
return self.large_heap[0]