Software Engineer's Blog

295. Find Median from Data Stream

295. Find Median from Data Stream

Sorting the whole stream on every query is the obvious answer, and it’s also the one that times out. The real trick is to keep the numbers half-sorted around the middle, and a pair of heaps does exactly that. It’s the headline example in the heap and priority-queue pattern — once the two-heap balance clicks here, running medians and “kth largest so far” problems feel like the same move.

The problem

You build a MedianFinder that ingests integers one at a time via addNum, and at any point findMedian must return the median of everything seen so far. Odd count gives the middle element; even count gives the average of the two middle ones. (Full statement on LeetCode.)

Feed it 1, then 2, and findMedian returns 1.5. Add 3 and it returns 2.0. The catch is scale: with tens of thousands of interleaved adds and queries, you can’t afford to re-sort each time.

Intuition: split the stream at the median with two heaps

Picture the sorted stream cut down the middle into a lower half and an upper half. The median only ever depends on the boundary between those halves — the largest of the lower half, the smallest of the upper half, or the average of the two. You never need the full ordering, just fast access to that seam.

That’s precisely what heaps give you. Keep the lower half in a max-heap, so its top is the biggest small value. Keep the upper half in a min-heap, so its top is the smallest large value. If you hold two invariants:

  • every element in the max-heap is \le every element in the min-heap, and
  • the two sizes differ by at most one (with the max-heap allowed to hold the extra),

then the median is a peek away. Equal sizes means the answer straddles both tops; when the max-heap is larger by one, its top is the median.

Adding a number is where the balancing happens. Push it onto the max-heap, then immediately move that heap’s top over to the min-heap. That single hop guarantees the “everything on the left \le everything on the right” invariant, because the value shipped across is the current largest of the lower half. It can leave the min-heap one too big, so if that happens, hop its top back. Each offer/poll is O(logn)O(\log n); the peek at query time is O(1)O(1).

Solution

import java.util.Collections;
import java.util.PriorityQueue;

class MedianFinder {

    // Lower half — top is the largest of the small numbers.
    private final PriorityQueue<Integer> maxHeap =
            new PriorityQueue<>(Collections.reverseOrder());
    // Upper half — top is the smallest of the large numbers.
    private final PriorityQueue<Integer> minHeap = new PriorityQueue<>();

    public void addNum(int num) {
        maxHeap.offer(num);
        // Ship the lower half's max across so left <= right holds.
        minHeap.offer(maxHeap.poll());
        // Rebalance: let the max-heap carry the extra element on odd counts.
        if (maxHeap.size() < minHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }

    public double findMedian() {
        // Odd total: the extra element on the max-heap is the middle.
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.peek();
        }
        // Even total: average the two values flanking the middle.
        return (maxHeap.peek() + minHeap.peek()) / 2.0;
    }
}

The / 2.0 matters — divide by an int 2 and (1 + 2) / 2 silently truncates to 1, not 1.5. Routing every insert through the max-heap first keeps the flow uniform, so there’s no special case for the first element or for which heap is currently larger.

Complexity

OperationTimeSpace
addNumO(logn)O(\log n)O(n)O(n)
findMedianO(1)O(1)

Total space is O(n)O(n) since both heaps together hold every element seen. The win over a re-sort-per-query approach is stark: O(logn)O(\log n) inserts versus O(nlogn)O(n \log n).

In an interview

Say the framing before you write code: “I’ll keep the smaller half in a max-heap and the larger half in a min-heap, balanced so the median sits at the tops.” That one sentence signals you’ve found the structure, which is what earns the checkmark.

The bug they wait for is dividing by 2 instead of 2.0 on the even case — an integer-truncation trap that passes small tests and fails on [1, 2]. Two more worth naming out loud: the size invariant (decide up front that the max-heap holds the extra, so findMedian has no ambiguity), and the follow-up. If the stream is bounded to [0, 100], mention that a 101-bucket counting array turns both operations near-O(1)O(1) — a good sign you see beyond the general solution. The same balance-two-structures idea recurs throughout the heap and priority-queue pattern.

References