Problem 34 · Heaps

Find Median from Data Stream

HardClassic DSA & systems

01Problem

Design a structure supporting add_num(x) and find_median() over a stream of integers.

02Key insight

Two heaps: a max-heap for the lower half and a min-heap for the upper half, kept balanced within one element; the median is read from the tops in O(1).

03Solution

solution.pypython19 lines

import heapq

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

    def add_num(self, num: int) -> None:
        heapq.heappush(self.low, -num)
        heapq.heappush(self.high, -heapq.heappop(self.low))
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))

    def find_median(self) -> float:
        if not self.low:
            raise ValueError("no elements")
        if len(self.low) > len(self.high):
            return float(-self.low[0])
        return (-self.low[0] + self.high[0]) / 2

Try it first.

TimeO(log n) add, O(1) median

SpaceO(n)

tests.py · 3 passing

tests.pypython8 lines

mf = MedianFinder()
mf.add_num(1); mf.add_num(2)
assert mf.find_median() == 1.5
mf.add_num(3)
assert mf.find_median() == 2.0
for x in [-5, -5, 100]:
    mf.add_num(x)
assert mf.find_median() == 1.5

04Staff-level follow-up

Unbounded stream across many machines? Exact medians need all data; use mergeable quantile sketches (t-digest, KLL) with bounded error, or bucket counts if the value range is small.