Find Median from Data Stream

Share

Design a data structure that supports:

Example
MedianFinder mf = new MedianFinder();
mf.addNum(1);
mf.addNum(2);
mf.findMedian(); → 1.5
mf.addNum(3);
mf.findMedian(); → 2.0
Constraints:

Contents

Approach 1: Sorted List (Brute Force)

Insert each number into a sorted list. Median is the middle element. O(n) insert, O(1) median.

Java
class MedianFinder {
    private List<Integer> sorted = new ArrayList<>();

    public void addNum(int num) {
        // Binary search for correct insertion point — O(log n) search, O(n) insert
        int pos = Collections.binarySearch(sorted, num);
        if (pos < 0) pos = -pos - 1;
        sorted.add(pos, num); // O(n) shift
    }

    public double findMedian() {
        int n = sorted.size();
        if (n % 2 == 1) return sorted.get(n / 2);
        return (sorted.get(n / 2 - 1) + sorted.get(n / 2)) / 2.0;
    }
}

Approach 2: Two Heaps (Optimal)

Maintain two heaps that partition all numbers into two equal halves:

Invariant: lowerHalf.size() == upperHalf.size() (even count) or lowerHalf.size() == upperHalf.size() + 1 (odd count. Extra element lives in lowerHalf).

Java
class MedianFinder {
    // Max-heap for lower half
    private PriorityQueue<Integer> lowerHalf = new PriorityQueue<>(Collections.reverseOrder());
    // Min-heap for upper half
    private PriorityQueue<Integer> upperHalf = new PriorityQueue<>();

    public void addNum(int num) {
        // Always add to lowerHalf first
        lowerHalf.offer(num);

        // Ensure all lowerHalf elements <= all upperHalf elements
        if (!upperHalf.isEmpty() && lowerHalf.peek() > upperHalf.peek()) {
            upperHalf.offer(lowerHalf.poll());
        }

        // Balance sizes: lowerHalf can have at most 1 extra element
        if (lowerHalf.size() > upperHalf.size() + 1) {
            upperHalf.offer(lowerHalf.poll());
        } else if (upperHalf.size() > lowerHalf.size()) {
            lowerHalf.offer(upperHalf.poll());
        }
    }

    public double findMedian() {
        if (lowerHalf.size() == upperHalf.size()) {
            // Even number of elements — average the two middles
            return (lowerHalf.peek() + upperHalf.peek()) / 2.0;
        }
        // Odd — extra element in lowerHalf is the median
        return lowerHalf.peek();
    }
}
The two-heap invariant: max of lowerHalf <= min of upperHalf. When violated (after adding a large number to lowerHalf), move the max of lowerHalf to upperHalf.

Complexity Analysis
ApproachaddNumfindMedianSpace
Sorted ListO(n)O(1)O(n)
Two HeapsO(log n)O(1)O(n)