Find Median from Data Stream

hardHeap / Priority Queue27 stepsJava · Python · C++

Keep the smaller half of the numbers in a max-heap and the larger half in a min-heap, balanced so the max-heap has the same size or one more. The median is the max-heap's top, or the average of both tops.

▶ Watch the animated trace

Traced input: addNum(5), addNum(15), findMedian(), addNum(1), findMedian(), addNum(3), findMedian() → [null, null, null, 10.0, null, 5.0, null, 4.0]. Original problem: neetcode.io/problems/find-median-in-a-data-stream

Complexity

BestO(log n)

AverageO(log n)

WorstO(log n)

SpaceO(n)

addNum is O(log n); findMedian is O(1).

Key idea

low (max-heap) holds the smaller half; high (min-heap) holds the larger half.
Every value in low is ≤ every value in high.
Add to low, move its top to high, then move back if high got bigger.
Odd count: median is low.peek(). Even: average the two tops.

Common mistakes

Integer division when averaging; use 2.0.
Letting the heaps drift more than one apart.
Sorting on every add is O(n log n) per call.

Solution

Java

class MedianFinder {
    private final PriorityQueue<Integer> low = new PriorityQueue<>(Collections.reverseOrder());
    private final PriorityQueue<Integer> high = new PriorityQueue<>();

    void addNum(int num) {
        low.offer(num);
        high.offer(low.poll());
        if (high.size() > low.size()) low.offer(high.poll());
    }

    double findMedian() {
        double m = low.size() > high.size() ? low.peek() : (low.peek() + high.peek()) / 2.0;
        return m;
    }
}

Python

class MedianFinder:
    def __init__(self):
        self.low = []   # max-heap (stored negated): smaller half
        self.high = []  # min-heap: larger half

    def addNum(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 findMedian(self) -> float:
        if len(self.low) > len(self.high): return float(-self.low[0])
        return (-self.low[0] + self.high[0]) / 2.0

C++

class MedianFinder {
    priority_queue<int> low;                             // smaller half, max on top
    priority_queue<int, vector<int>, greater<int>> high;  // larger half, min on top

public:
    void addNum(int num) {
        low.push(num);
        high.push(low.top()); low.pop();
        if (high.size() > low.size()) { low.push(high.top()); high.pop(); }
    }

    double findMedian() {
        if (low.size() > high.size()) return low.top();
        return (low.top() + high.top()) / 2.0;
    }
};

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
Single number The max-heap holds the only value.	`addNum(7), findMedian()`	`[null, null, 7.0]`
Negative numbers Average of -2 and -1.	`addNum(-1), addNum(-2), findMedian()`	`[null, null, null, -1.5]`
Descending input Rebalancing keeps the middle value on top of `low`.	`addNum(3), addNum(2), addNum(1), findMedian()`	`[null, null, null, null, 2.0]`