Find Median from Data Stream: The Two Heaps Pattern Explained

Numbers arrive one at a time, and after each one you must report the median of everything so far. Sorting after every arrival works and is far too slow. The standard answer uses two heaps that face each other, so the middle of the data is always sitting on top of one of them. It is the entry point to a whole family of "two heaps" problems, and it is short enough to write from memory.

The problem it solves

Design a structure with two operations: add_num(x) records a value, and find_median() returns the median of all values recorded so far. For an odd count the median is the middle value; for an even count it is the average of the two middle values. For the stream 5, 15, 1, 3 the medians are 5, 10, 5 and 4.

The naive versions each pay somewhere. Re-sorting on every read is O(n log n) per read. Keeping a sorted list with binary insertion makes the search O(log n), but shifting the list to insert is O(n). With a million values arriving, that shift is what hurts.

The intuition

You never need the whole order, only the boundary between the lower half and the upper half. Heaps are good at exactly one thing: showing you the largest or the smallest value in O(1).

• low is a max-heap of the smaller half. Its top is the largest of the small values.
• high is a min-heap of the larger half. Its top is the smallest of the large values.

Keep two invariants: every value in low is at most every value in high, and low has either the same number of values as high or one more. Then the median is low's top when the count is odd, and the average of the two tops when it is even.

Inserting without deciding which side a value belongs to is the neat part. Every value enters low, and low immediately hands its largest value to high. That round trip guarantees the ordering invariant: whatever crosses is the largest of the low side, so it belongs on the high side. If high is now bigger, it hands its smallest back. Each step is a heap push or pop, O(log n).

Watch it run

The animation draws low and high as two trees whose roots face each other. 5 enters low, crosses to high, and comes back because high became bigger: median 5. 15 enters low and crosses to high; the halves are even, so the median is (5 + 15) / 2 = 10. 1 enters low, 5 crosses, and 5 comes back: median 5. 3 enters low, 5 crosses, and the halves are even again: (3 + 5) / 2 = 4. Nothing was ever re-sorted.

The code

Python's heapq is a min-heap, so low stores negated values:

import heapq

class MedianFinder:
    def __init__(self):
        self.low = []                          # max-heap of the smaller half, stored negated
        self.high = []                         # min-heap of the larger half

    def add_num(self, x):
        heapq.heappush(self.low, -x)                         # everything enters low
        heapq.heappush(self.high, -heapq.heappop(self.low))  # low's largest crosses over
        if len(self.high) > len(self.low):                   # low keeps the extra value
            heapq.heappush(self.low, -heapq.heappop(self.high))

    def find_median(self):
        if len(self.low) > len(self.high):
            return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2

mf = MedianFinder()
out = []
for x in [5, 15, 1, 3]:
    mf.add_num(x)
    out.append(mf.find_median())
print(out)                                     # [5, 10.0, 5, 4.0]
print(sorted(-v for v in mf.low), sorted(mf.high))   # [1, 3] [5, 15]

mf = MedianFinder()
for x in [2, 2, 2, -7]:
    mf.add_num(x)
print(mf.find_median())                        # 2.0

Against a sorted list kept with bisect.insort, on 500 random streams with duplicates and negative numbers, checking the two invariants after every insert:

import bisect, random

class SortedListMedian:
    def __init__(self):
        self.values = []

    def add_num(self, x):
        bisect.insort(self.values, x)          # O(log n) search, O(n) shift

    def find_median(self):
        v, n = self.values, len(self.values)
        return v[n // 2] if n % 2 else (v[n // 2 - 1] + v[n // 2]) / 2

random.seed(16)
ok = True
for _ in range(500):
    a, b = MedianFinder(), SortedListMedian()
    for _ in range(random.randint(1, 60)):
        x = random.randint(-50, 50)
        a.add_num(x)
        b.add_num(x)
        ok &= a.find_median() == b.find_median()
        ok &= len(a.low) - len(a.high) in (0, 1)
        ok &= not a.high or -a.low[0] <= a.high[0]
print(ok)                                      # True

The complexity

• Re-sort on each read: O(n log n) per read.
• Sorted list with binary insertion: O(n) per insert because of the shift, O(1) per read.
• Two heaps: at most five heap operations per insert, so O(log n); the read looks at two tops, O(1). Space is O(n), since every value is kept.

Where it goes wrong

• Forgetting to negate on the way out. low holds negatives; -self.low[0] is its largest real value. Missing one minus sign gives a median with the wrong sign.
• Balancing sizes but not order. Pushing a value into whichever heap is smaller keeps the sizes level but can put a large value in low. The round trip through low is what keeps every low value at most every high value.
• Integer division for the even case. // truncates (3 + 4) // 2 to 3. Use /, or the language's floating-point division.
• Reading from an empty structure. find_median before any add_num touches self.low[0]; decide whether that raises or returns nothing.

When it shows up in interviews

It shows up as a hard-labelled question that is really a data structure design question: can you pick the right structure for "keep the middle visible"? The follow-ups are the interesting part. If every value is between 0 and 100, a count array of 101 buckets answers in constant time. If only the last k values count, you need removal from a heap, which is the sliding window median. Outside interviews, a running median is how dashboards report a typical latency or price without the average being dragged around by outliers.

How to say it in an interview

"I keep the smaller half in a max-heap and the larger half in a min-heap, with the low side allowed one extra value. Each new number goes into the low heap, then the low heap's largest moves to the high heap, and if the high heap is now bigger its smallest moves back. That keeps every low value at most every high value. The median is the low top, or the average of both tops. Insert is O(log n), the read is O(1)."

The removal problem is solved in sliding window median with two heaps, and the heap itself is introduced in heap basics.