3 min readRishi

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

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.