4.17.7 — Find Median from Data Stream

LeetCode 295 · Hard · ★ Blind 75

The problem

Numbers arrive one at a time. After any number of them, return the median of everything seen so far.

addNum(1); addNum(2)
findMedian()  →  1.5
addNum(3)
findMedian()  →  2.0

Up to 50,000 calls.

The pattern

Keeping a sorted array gives O(1) median but O(n) insertion, because of the shifting. Keeping an unsorted array gives O(1) insertion but O(n \log n) median. Neither is good enough.

The observation that fixes it:

You do not need the data sorted. You only need the two values in the middle.

So split the numbers into a small half and a large half, keeping them the same size (or the small half one bigger). Then:

• The largest of the small half, and the smallest of the large half, are the two middle values.
• A max-heap gives you the largest of the small half in O(1).
• A min-heap gives you the smallest of the large half in O(1).

Two heaps facing each other, with the median sitting between their roots.

        small half              large half
      (max-heap)                (min-heap)
    [1, 2, 3]                   [4, 5, 6]
          ↑                      ↑
        root = 3               root = 4
                median = 3.5

The two invariants

Everything in the code exists to maintain these:

1. Every value in the small half ≤ every value in the large half.
2. The sizes differ by at most 1, and if they differ, the small half is the larger one.

Invariant 2's tie-break is a choice, not a law. Making the small half the bigger one means an odd count has its median at the small half's root, so findMedian is a one-line check.

Maintaining them on insert

Two steps, and doing them in this order removes every special case:

Step 1 — push onto the small half, then move its largest across to the large half.

Pushing blindly could break invariant 1, so immediately hand the small half's new maximum to the large half. After this, invariant 1 always holds, because the value that moved was the largest thing on the small side.

Step 2 — if the large half is now bigger, move its smallest back.

That restores invariant 2.

No if about which heap the value belongs in. Every insert follows the same three heap operations.

The solution

Python

import heapq

class MedianFinder:
    def __init__(self):
        self.small = []      # max-heap (negated), the lower half
        self.large = []      # min-heap, the upper half

    def addNum(self, num: int) -> None:
        heapq.heappush(self.small, -num)                      # always goes here first
        heapq.heappush(self.large, -heapq.heappop(self.small))  # hand over the largest

        if len(self.large) > len(self.small):                 # rebalance
            heapq.heappush(self.small, -heapq.heappop(self.large))

    def findMedian(self) -> float:
        if len(self.small) > len(self.large):
            return float(-self.small[0])                      # odd count
        return (-self.small[0] + self.large[0]) / 2.0         # even count

TypeScript

class MedianFinder {
  private small = new MaxHeap<number>();   // lower half
  private large = new MinHeap<number>();   // upper half

  addNum(num: number): void {
    this.small.push(num);
    this.large.push(this.small.pop()!);
    if (this.large.size > this.small.size) {
      this.small.push(this.large.pop()!);
    }
  }

  findMedian(): number {
    if (this.small.size > this.large.size) return this.small.peek()!;
    return (this.small.peek()! + this.large.peek()!) / 2;
  }
}

Python's heapq is min-only, so the small half stores negated values. -self.small[0] reads the true maximum back. Negating on the way in and out is the standard trick, and forgetting one of the two is the standard bug.

Read the addNum body as one sentence: push it onto the small side, immediately give the small side's largest to the large side, and if the large side has grown too big, give its smallest back. Three lines, no branches, invariants always restored.

Trace

addNum(1): push to small → small = [1]. Move largest across → small = [], large = [1]. Large is bigger → move back → small = [1], large = []. Median is 1.

addNum(2): push to small → small = [1, 2]. Move largest (2) across → small = [1], large = [2]. Sizes equal. Median is (1 + 2) / 2 = 1.5 ✓.

addNum(3): push to small → small = [1, 3]. Move 3 across → small = [1], large = [2, 3]. Large is bigger → move 2 back → small = [1, 2], large = [3]. Median is 2 ✓.

Complexity

addNum is O(\log n) — three heap operations. findMedian is O(1).

Space is O(n), which is unavoidable when the median of everything so far may be asked for at any moment.

The follow-ups, and they are asked

"What if all the numbers are between 0 and 100?"

A bounded range means counting. Keep a 101-slot array of counts and a running total. findMedian scans the buckets accumulating counts until it passes the halfway point — O(101), a constant.

addNum becomes O(1) and space becomes O(1). The same "bounded values become array indices" realisation as 4.4.5, and this follow-up is asked specifically to see whether you have it.

"What if 99% of the numbers are between 0 and 100?"

Count the common range in buckets, and keep the rare outliers in two small sorted lists — one below the range, one above. The median lookup checks the outlier counts first, then falls into the bucket scan. Almost all the data costs O(1), and the exceptions are few enough to handle exactly.

Where this goes next

• Sliding Window Median — the median of a moving window, which needs deletion from the middle of a heap. Standard heaps cannot delete efficiently, so you use lazy deletion with a "to remove" map, or a balanced BST / order-statistic tree (4.14.12).
• IPO / Maximum Capital — two heaps again, one for what is affordable and one for what is not yet.
• Percentile monitoring — real systems track p50, p95 and p99 over a stream and cannot store everything, so they use approximate structures like t-digest. That is the honest production answer.

The rule: two heaps facing each other give you the middle of a stream in O(\log n) per insert. Whenever a problem needs a boundary between two groups that must stay balanced, reach for it.

What the interviewer will push on

"Why two heaps?" You need only the two middle values, and each heap hands you one of them in O(1).

"State your invariants." Small ≤ large elementwise, and sizes within 1 with the small half favoured.

"Why push to small first, always?" So the insert has no branch. The immediate hand-over restores the ordering invariant, and the rebalance restores the size invariant.

"What if the values are bounded?" Counting buckets, O(1) per add.

"What about 99% in a narrow range?" Buckets plus small sorted lists for the outliers.

"Can you delete a value?" Not directly from a heap. Lazy deletion or a different structure.

One thing to volunteer: state the two invariants before writing any code, then write the three lines and show they restore both. This problem is short but the reasoning is what is being marked.

Next: 4.18 is the enumeration group — every subset, permutation and arrangement — where the code is always the same four parts and all the difficulty is in the pruning.