Find Median from Data Stream — Two-Heap Approach

Problem Statement#

Design a data structure that supports two operations: addNum(num) adds a number from an integer data stream; findMedian() returns the median of all numbers added so far. For even count, the median is the average of the two middle elements.

Constraints:

• -10^5 <= num <= 10^5
• At most 5 * 10^4 calls to addNum and findMedian
• At least one number is added before findMedian is called

Input:  addNum(1), addNum(2), findMedian(), addNum(3), findMedian()
Output: 1.5, 2.0

Input:  addNum(6), addNum(10), addNum(2), addNum(6), findMedian()
Output: 6.0

Why This Problem Matters#

Running medians appear in analytics pipelines, A/B test result streaming, stock price monitoring, and any system that needs real-time statistical summaries over unbounded data. Google's streaming analytics platform, Apache Flink, and Amazon Kinesis Data Analytics all implement variants of this algorithm for rolling quantile computation.

This is a canonical hard problem in FAANG interviews because it requires combining two heaps in a non-obvious way and maintaining a careful invariant across every insert. The two-heap approach achieves O(log N) per insert and O(1) median—optimal for a streaming context where you cannot store all values.

Understanding this problem also provides intuition for the general k-th order statistic in a stream, percentile tracking, and the P99 latency computation used in SLA monitoring at every major technology company.

The Core Insight#

Partition numbers into two halves: lo (max-heap containing the lower half) and hi (min-heap containing the upper half). Maintain the invariant: len(lo) >= len(hi) and every element in lo is less than or equal to every element in hi.

After each addNum, push to lo then transfer lo's max to hi. If hi becomes larger than lo, transfer hi's min back to lo. This rebalancing ensures the invariant holds.

Median: if total count is odd, return lo's max. If even, return the average of lo's max and hi's min. In Python, negate values to simulate a max-heap with heapq (which is a min-heap).

Visual Dry Run#

Solution (Optimal)#

import heapq
 
class MedianFinder:
    def __init__(self):
        self.lo = []   # max-heap (negate values)
        self.hi = []   # min-heap
 
    def addNum(self, num: int) -> None:
        heapq.heappush(self.lo, -num)
        heapq.heappush(self.hi, -heapq.heappop(self.lo))
        if len(self.hi) > len(self.lo):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))
 
    def findMedian(self) -> float:
        if len(self.lo) > len(self.hi):
            return float(-self.lo[0])
        return (-self.lo[0] + self.hi[0]) / 2.0

class MedianFinder {
    constructor() {
        this.lo = [];  // max-heap stored as negated min-heap
        this.hi = [];  // min-heap
    }
 
    addNum(n) {
        // Push to lo (max-heap via negation), then balance
        this._pushMin(this.lo, -n);
        this._pushMin(this.hi, -this._popMin(this.lo));
        if (this.hi.length > this.lo.length) {
            this._pushMin(this.lo, -this._popMin(this.hi));
        }
    }
 
    findMedian() {
        if (this.lo.length > this.hi.length) return -this.lo[0];
        return (-this.lo[0] + this.hi[0]) / 2;
    }
 
    _pushMin(h, v) {
        h.push(v);
        let i = h.length - 1;
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (h[p] > h[i]) { [h[p], h[i]] = [h[i], h[p]]; i = p; }
            else break;
        }
    }
 
    _popMin(h) {
        [h[0], h[h.length - 1]] = [h[h.length - 1], h[0]];
        const v = h.pop();
        let i = 0;
        while (true) {
            let m = i, l = 2*i+1, r = 2*i+2;
            if (l < h.length && h[l] < h[m]) m = l;
            if (r < h.length && h[r] < h[m]) m = r;
            if (m === i) break;
            [h[i], h[m]] = [h[m], h[i]]; i = m;
        }
        return v;
    }
}

Time: O(log N) for addNum — two heap pushes and at most two pops each
Space: O(N) — all elements stored across the two heaps

Common Mistakes#

• Pushing to hi first instead of lo—the invariant requires all elements flow through lo first so that the max of lo is correctly propagated to hi
• Forgetting to negate values in Python: heapq is a min-heap; push -num to simulate a max-heap for lo
• Off-by-one in the size check: rebalance when len(hi) > len(lo), not when they are equal—the invariant allows lo to have one extra element
• Returning int from findMedian instead of float—the average of two integers must use floating-point division
• Not handling the single-element case: with one element in lo and nothing in hi, median is lo[0] directly

Interview Tips#

• Draw the two heaps explicitly and label them "lower half max-heap" and "upper half min-heap"
• State the invariant before writing any code: "lo has all smaller-half values, hi has larger-half values; lo can have one more element than hi"
• The three-step addNum pattern is canonical: push to lo, transfer lo's max to hi, rebalance if needed
• In Python, explain the negation trick for max-heap simulation upfront—interviewers who don't know Python may be confused

Follow-up Questions#

• How would you find the k-th smallest element in a stream instead of the median? (Use a min-heap of size k; the top is always the k-th largest seen so far)
• How would you handle a sliding window median (last N elements only)? (Two heaps plus a lazy deletion set to remove expired elements)
• What if the stream has a skewed distribution with mostly small numbers? (The invariant self-adjusts; no additional logic needed)
• How would you compute the 99th percentile latency in a stream? (Maintain two heaps with a 99/1 size ratio instead of 50/50)
• How would you make this thread-safe for concurrent addNum calls? (Lock both heaps together atomically; or use a lock-free priority queue)

Key Takeaways#

• lo is a max-heap (lower half) and hi is a min-heap (upper half); maintain len(lo) >= len(hi)
• addNum always pushes to lo first, transfers lo's max to hi, then rebalances if hi grows larger
• In Python, simulate a max-heap by negating all values pushed to lo
• findMedian returns lo's max for odd total count, or the average of lo's max and hi's min for even count
• This achieves O(log N) insert and O(1) median—optimal for an unbounded data stream
• The same two-heap partition generalises to any percentile tracking: change the size ratio from 50/50 to p/(1-p)
• LeetCode 295 is the canonical version; Google and Microsoft use variants with sliding window constraints