#295HardHeaps & top-k

Find Median from Data Stream

The tell: A running order statistic over an unbounded stream.

How you get there

1Keep a max-heap of the lower half and a min-heap of the upper half.
2Push, then rebalance so the sizes differ by at most one.
3The median is a heap root, or the mean of both roots.

Solution

Python

import heapq

class MedianFinder:
    def __init__(self) -> None:
        self._low: list[int] = []    # max-heap via negation
        self._high: list[int] = []   # min-heap

    def addNum(self, num: int) -> None:
        heapq.heappush(self._low, -num)
        heapq.heappush(self._high, -heapq.heappop(self._low))   # keep ordering invariant
        if len(self._high) > len(self._low):
            heapq.heappush(self._low, -heapq.heappop(self._high))

    def findMedian(self) -> float:
        if len(self._low) > len(self._high):
            return float(-self._low[0])
        return (-self._low[0] + self._high[0]) / 2

Time: O(log n) per insertion, O(1) per query
Space: O(n)

What goes wrong

Balancing sizes without first moving an element across breaks the ordering invariant between the heaps.

Test yourself on DSA Next: Kth Largest Element in an Array →