Find Median from Data Stream

Find Median from Data Stream is the classic interview test of whether you reach for the right data structure. Numbers arrive one at a time, and after each one you must be able to report the median instantly. The trick: split the data into two heaps that meet in the middle.

The slow way first

The naive idea: keep all the numbers in a list, and on every find_median sort the list and pick the middle. Sorting is O(n log n) per query — and you might query after every insert, so this is painfully slow on a long stream.

A small improvement is to keep the list sorted and insert each new number into place. But shifting elements to make room is O(n) per insert. We want both operations fast.

The idea: two heaps meeting in the middle

Keep the lower half of the numbers in a max-heap and the upper half in a min-heap. Then:

• the largest of the lower half is at the max-heap's top,
• the smallest of the upper half is at the min-heap's top,

and those two tops are exactly the middle of the sorted data. If we keep the two heaps balanced (sizes differing by at most one), the median is just a peek at the top(s) — O(1).

Python only has a min-heap (heapq), so we fake a max-heap by storing negated values: the most-negative entry (the true max) floats to the top.

Walk through it

Step through the animation as 5, 3, 8, 2 stream in. Each number first lands on a side (lower if it is <= the current max of the lower half, otherwise upper). Then a rebalance moves one element across if a heap grew too big. Watch how the max-heap top and min-heap top hug the middle, so the median label updates in O(1) after every insert.

Pseudocode

lo = max-heap   # lower half
hi = min-heap   # upper half
 
add_num(x):
    if lo is empty or x <= top(lo):
        push x onto lo
    else:
        push x onto hi
    # rebalance so |lo| and |hi| differ by at most 1, lo never smaller
    if size(lo) > size(hi) + 1: move top(lo) -> hi
    elif size(hi) > size(lo):   move top(hi) -> lo
 
find_median():
    if size(lo) > size(hi): return top(lo)
    return (top(lo) + top(hi)) / 2

The Python solution

import heapq
 
class MedianFinder:
    def __init__(self):
        self.lo = []   # max-heap (store negatives)
        self.hi = []   # min-heap
 
    def add_num(self, num):
        if not self.lo or num <= -self.lo[0]:
            heapq.heappush(self.lo, -num)
        else:
            heapq.heappush(self.hi, num)
        # rebalance so sizes differ by at most 1
        if len(self.lo) > len(self.hi) + 1:
            heapq.heappush(self.hi, -heapq.heappop(self.lo))
        elif len(self.hi) > len(self.lo):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))
 
    def find_median(self):
        if len(self.lo) > len(self.hi):
            return -self.lo[0]
        return (-self.lo[0] + self.hi[0]) / 2

• self.lo is a max-heap simulated by pushing -num, so -self.lo[0] is its true maximum.
• The if num <= -self.lo[0] check routes a number to the correct half before rebalancing.
• The two rebalance branches move a single element across whenever a heap gets too big, keeping lo either equal to or exactly one larger than hi.
• find_median is O(1): if lo is bigger the median is its top; otherwise the count is even and we average the two tops.

Complexity

We hold every number across the two heaps, so space is O(n), but each insert is only O(log n) and every median read is O(1).

When this pattern shows up

Practice