Find Median from Data Stream
The Prompt
The median is the middle value in an ordered integer list. Design a data structure that supports the following two operations: `addNum` which adds an integer number from the data stream to the data structure, and `findMedian` which returns the median of all elements so far.
Understanding the Problem
The median needs the middle of a sorted order, but the stream never stops: keeping a sorted array makes addNum O(n) per insert, and sorting on every findMedian is worse. You do not actually need the whole order — only the one or two elements at the boundary between the smaller half and the larger half.
So maintain that boundary directly with two heaps: a max-heap holding the smaller half (its root is the largest small number) and a min-heap holding the larger half (its root is the smallest large number). The median is always read off the roots in O(1).
The Interview Flow
Interviewer
How would you design a data structure to efficiently find the median of a stream of numbers?
Candidate
Keeping the full list sorted and inserting new numbers would be O(n) for each `addNum`. That's too slow. I need a way to access the middle elements quickly.
Interviewer
What data structure can help with that?
Candidate
I can use two heaps. A max-heap to store the smaller half of the numbers, and a min-heap to store the larger half. The max-heap will have the largest of the small numbers at its root, and the min-heap will have the smallest of the large numbers at its root. These two roots are the middle elements.
Interviewer
How do you maintain this structure when adding a new number?
Candidate
When a new number comes in, I add it to the max-heap (for the smaller half). Then, to maintain balance, I move the largest element from the max-heap to the min-heap. After this, I need to ensure the heaps are balanced in size. The max-heap can have at most one more element than the min-heap.
Interviewer
How do you balance the sizes?
Candidate
If the min-heap becomes larger than the max-heap, I move its smallest element (the root) to the max-heap. This keeps the sizes either equal or with the max-heap being one larger.
Interviewer
And how do you find the median?
Candidate
If the total number of elements is odd, the median is simply the root of the max-heap (which is the larger heap). If the total is even, the median is the average of the roots of both heaps.
Interviewer
This is a very clever use of two heaps. The `addNum` operation would be O(log n) and `findMedian` would be O(1). Perfect. Please describe the code structure.
Why do two balanced heaps always expose the median?
Two invariants: (1) every element in the max-heap ≤ every element in the min-heap, and (2) their sizes differ by at most one, with the max-heap allowed the extra element. Together they force the roots to be the middle of the sorted order: with an odd count the median is the max-heap root; with an even count it is the average of the two roots.
Each addNum does at most a few heap pushes and pops — O(log n) — and findMedian is O(1). That is the trade: you give up a fully sorted structure (unneeded) to make the one query you care about constant-time, versus O(n) inserts for a sorted list or O(n log n) re-sorting per query.
Two Heaps Approach
- Initialize two heaps: a `small` heap (max-heap) and a `large` heap (min-heap).
- **For `addNum(num)`:**
- Add `num` to the `small` (max-heap). In many languages, you add `-num` to a min-heap to simulate a max-heap.
- Balance by moving the largest element from `small` to `large`. Pop from `small` and push to `large`.
- Re-balance sizes. If `len(large) > len(small)`, move the smallest element from `large` to `small`. Pop from `large` and push to `small`.
- **For `findMedian()`:**
- If `len(small) > len(large)`, the total count is odd. The median is the root of `small`.
- If `len(small) == len(large)`, the total count is even. The median is the average of the roots of `small` and `large`.
Try it yourself
Write your solution and run it against 2 test cases.
JavaScript, TypeScript & Python run sandboxed in your browser; other languages run on the execution server. Your code is saved locally as you type.
Final Solution
// JavaScript doesn't have a native heap, so this requires a library
// or a custom implementation. The following is a conceptual example.
class MedianFinder {
constructor() {
// max heap
this.small = new MaxPriorityQueue();
// min heap
this.large = new MinPriorityQueue();
}
addNum(num) {
this.small.enqueue(num);
if (this.small.size() > 0 && this.large.size() > 0 && this.small.front().element > this.large.front().element) {
this.large.enqueue(this.small.dequeue().element);
}
if (this.small.size() > this.large.size() + 1) {
this.large.enqueue(this.small.dequeue().element);
}
if (this.large.size() > this.small.size() + 1) {
this.small.enqueue(this.large.dequeue().element);
}
}
findMedian() {
if (this.small.size() > this.large.size()) {
return this.small.front().element;
}
if (this.large.size() > this.small.size()) {
return this.large.front().element;
}
return (this.small.front().element + this.large.front().element) / 2;
}
}Explanation
Stream in 5, 15, 1, 3 and watch the two heaps split the numbers around the middle.
1addNum(5): with one element, the max-heap keeps it (it may hold one extra). Odd count, so median = max-heap root = 5.
2addNum(15): it enters the max-heap, but as the largest element it is moved to the min-heap. Even count: median = (5 + 15) / 2 = 10.
3addNum(1): after the push-across and rebalance, the small half is {1, 5} and the large half {15}. Odd count: median = max-heap root = 5.
4addNum(3): 3 joins the small half and 5 overflows to the large half — sizes 2 and 2. Median = (3 + 5) / 2 = 4, matching sorted [1, 3, 5, 15].
Complexity Analysis
TIME
O(log n) for add, O(1) for median
SPACE
O(n)
Finished working through this one?
Mark it complete to track it on your Data Structures path.