Find Median From a Data Stream With Two Heaps
"Find Median from Data Stream" is a Hard-level staple in Amazon, Google, Apple loops, and it is a textbook fit for the Top-K and Heaps pattern — one of the nine patterns that cover the bulk of what FAANG coding interviews actually test. Part 11 of 11 in the Top-K and Heaps arc.
The Problem
Design a structure supporting add_num(num) and find_median() over the numbers seen so far. The crown jewel of the heap pattern — the median needs both the largest of the lower half and the smallest of the upper half, and no single heap exposes both.
add(1), add(2) -> median 1.5
add(3) -> median 2.0
Recognizing the Top-K and Heaps Pattern
When a problem asks for the k largest, smallest, closest, or most frequent — or for repeated access to an extreme while data changes — a heap gives O(log n) insertion and O(1) access to the extreme. The signature trick: keep a bounded heap of size k with the opposite polarity (a min-heap to track the k largest), evicting the root on overflow. Python's heapq is a min-heap; negate values for max behavior.
A running median partitions the stream into halves around it: a max-heap holding the lower half and a min-heap the upper half put both boundary elements at the roots. Maintaining a size invariant (lower holds equal or one extra) makes the median a root read — O(1) query after O(log n) insert.
The Approach
Insert by pushing onto the low (max) heap, moving its root to the high heap, then rebalancing back if the high side grew larger — three heap operations, unconditionally correct regardless of where the value lands. Median: low's root for odd counts, mean of both roots for even.
The always-rotate insert avoids the branchy compare-first version people fumble live. Python note: the low half stores negated values — negate on every read.
Python Solution
import heapq
class MedianFinder:
"""Streaming median via balanced max-heap (low) and min-heap (high)."""
def __init__(self) -> None:
self.low: list[int] = [] # max-heap (negated), holds equal or one extra
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 len(self.low) > len(self.high):
return float(-self.low[0])
return (-self.low[0] + self.high[0]) / 2.0
Complexity
- Time: O(log n) insert, O(1) median — three bounded heap operations per insert
- Space: O(n) — every element lives in one of the heaps
Interview Tips and Follow-Ups
- State the two invariants before coding — size (low equals high or one more) and order (low's max at most high's min). The code is just their maintenance.
- Follow-up one: all elements in [0, 100] — counting array beats heaps. Follow-up two: sliding-window median — lazy deletion or indexed containers. Both are standard.
- This structure powers percentile monitoring in real systems (p50 dashboards) — a one-sentence systems tie-in lands well at Google and Amazon.
More Top-K and Heaps problems — and the other eight patterns — live in the Technical Interview category. Drill the pattern, not the problem: that is the entire thesis of this series.
Keep reading
K Closest Points to Origin Without Square Roots
Bounded max-heap on squared distance — monotone transforms are free. Python solution and complexity analysis for the heap and top-k interview pattern.
Find K Pairs With Smallest Sums: Frontier Expansion
Explore an implicit sorted matrix — seed one row, expand neighbors lazily. Python solution and complexity analysis for the heap and top-k interview pattern.
Kth Largest Element: Heap Versus Quickselect
The bounded min-heap of size k — and when quickselect beats it. Python solution and complexity analysis for the heap and top-k interview pattern.
Newsletter
New posts, straight to your inbox
One email per post. No spam, no tracking pixels, unsubscribe anytime.
Comments
- No comments yet. Be the first.