☕ DSA Java Articles

Topics › Heaps › Hard Problems

Find Median in a Stream

Medium TUF ID 446 SDE

LeetCode ↗ Java source ↗ Edit this article ↗

Problem

Integers arrive one at a time. Design a structure that supports two operations:

The median is the middle value of the sorted data. When the count is even there is no single middle value, so the median is the average of the two middle values.

Constraints (LeetCode): -10^5 ≤ num ≤ 10^5, at most 5 · 10^4 calls in total, and findMedian is only called after at least one addNum. The code here accepts the full int range and throws IllegalStateException if asked for the median of an empty stream.

Example 1

addNum(1); addNum(2)
findMedian() -> 1.5      // sorted [1, 2]: average of 1 and 2
addNum(3)
findMedian() -> 2.0      // sorted [1, 2, 3]: the middle value

Example 2 (median after every add)

add 5  -> 5.0            // [5]
add 15 -> 10.0           // [5, 15]
add 1  -> 5.0            // [1, 5, 15]
add 3  -> 4.0            // [1, 3, 5, 15]: (3 + 5) / 2

All three designs below implement the same small interface, so one test script drives them all:

interface MedianFinder {
    void addNum(int num);
    double findMedian();
}

Approach 1: Brute force, sort on every query

Intuition. The median is defined on sorted data, so the most direct design stores numbers as they come and sorts them whenever the median is requested. Adding is trivial. The weakness is obvious: if one number arrives between two queries, the second query re-sorts everything even though the data barely changed.

Algorithm.

  1. addNum: append num to a list.
  2. findMedian: copy the list into an array and sort it. With n elements, return a[n/2] when n is odd, otherwise the average of a[n/2 - 1] and a[n/2].
static class SortOnQuery implements MedianFinder {
    private final List<Integer> nums = new ArrayList<>();

    public void addNum(int num) {
        nums.add(num);
    }

    public double findMedian() {
        int n = nums.size();
        if (n == 0) throw new IllegalStateException("stream is empty");
        int[] a = new int[n];
        for (int i = 0; i < n; i++) a[i] = nums.get(i);
        Arrays.sort(a);
        if (n % 2 == 1) return a[n / 2];
        return ((long) a[n / 2 - 1] + a[n / 2]) / 2.0;      // long: the sum can overflow int
    }
}

Complexity. addNum is O(1) amortised. findMedian is O(n log n) for the sort. Space O(n). If every add is followed by a query, n operations cost O(n^2 log n) in total.

Approach 2: Better, keep the list sorted at all times

Intuition. Instead of sorting from scratch at every query, maintain the invariant “the list is always sorted”. A new number then only has to be placed in the right position, which binary search finds in O(log n). The median is read directly from the middle indices.

The catch is the insertion itself: ArrayList.add(index, value) has to shift every element after index one step to the right. So binary search reduces the comparisons, but the move cost is still linear. In practice the shift is a single fast memory copy, so this design is far quicker than approach 1, but its worst case is still O(n) per add.

Algorithm.

  1. addNum: binary search for the first index whose value is greater than num (inserting after equal values keeps the search simple), then insert there.
  2. findMedian: read the middle element, or average the two middle elements.
static class SortedInsert implements MedianFinder {
    private final List<Integer> sorted = new ArrayList<>();

    public void addNum(int num) {
        int lo = 0, hi = sorted.size();                     // find the first index holding a value > num
        while (lo < hi) {
            int mid = (lo + hi) >>> 1;
            if (sorted.get(mid) <= num) lo = mid + 1; else hi = mid;
        }
        sorted.add(lo, num);                                // shifts the tail one step right: O(n)
    }

    public double findMedian() {
        int n = sorted.size();
        if (n == 0) throw new IllegalStateException("stream is empty");
        if (n % 2 == 1) return sorted.get(n / 2);
        return ((long) sorted.get(n / 2 - 1) + sorted.get(n / 2)) / 2.0;
    }
}

Complexity. addNum is O(log n) comparisons plus O(n) shifting, so O(n). findMedian is O(1). Space O(n).

Approach 3: Optimal, two heaps

Intuition. Look at what the median actually depends on. Split the sorted data into a smaller half and a larger half. The median is determined entirely by the largest value of the smaller half and the smallest value of the larger half. The order inside each half is irrelevant. We are paying for a fully sorted list when all we need are the two values next to the boundary.

A heap is the structure that gives cheap access to one extreme value and nothing more:

We maintain two invariants:

  1. Order: every value in low is ≤ every value in high.
  2. Size: low.size() == high.size() or low.size() == high.size() + 1.

With both in place, the median is low.peek() when the total count is odd (the extra element sits in low) and the average of the two peeks when it is even.

Why insertion keeps the order invariant. If num ≤ low.peek(), then num is no larger than the biggest element of the small half, which is no larger than anything in high, so num may go into low. Otherwise num is larger than everything in low, so it may go into high. Rebalancing moves low’s maximum to high, or high’s minimum to low. Each moved value is exactly the one at the boundary, so the order invariant survives the move, and the size invariant is restored.

Algorithm.

  1. addNum: push num into low if low is empty or num ≤ low.peek(), otherwise into high.
  2. If low now has two more elements than high, move low.poll() into high. If high is now larger than low, move high.poll() into low.
  3. findMedian: if low is larger, return low.peek(); otherwise return the average of low.peek() and high.peek().
static class TwoHeaps implements MedianFinder {
    private final PriorityQueue<Integer> low = new PriorityQueue<>(Collections.reverseOrder()); // smaller half, max on top
    private final PriorityQueue<Integer> high = new PriorityQueue<>();                          // larger half, min on top

    public void addNum(int num) {
        if (low.isEmpty() || num <= low.peek()) low.offer(num); else high.offer(num);
        if (low.size() > high.size() + 1) high.offer(low.poll());     // rebalance: low may hold one extra
        else if (high.size() > low.size()) low.offer(high.poll());
    }

    public double findMedian() {
        if (low.isEmpty()) throw new IllegalStateException("stream is empty");
        if (low.size() > high.size()) return low.peek();
        return ((long) low.peek() + high.peek()) / 2.0;
    }
}

Dry run on Example 2:

add goes into rebalance low (max-heap) high (min-heap) median
5 low (empty) none {5} {} 5.0
15 high (15 > 5) none {5} {15} 10.0
1 low (1 ≤ 5) none {5, 1} {15} 5.0
3 low (3 ≤ 5) low is 2 bigger: move 5 up {3, 1} {5, 15} (3 + 5) / 2 = 4.0

Complexity. addNum is O(log n): at most three heap operations. findMedian is O(1): two peeks. Space O(n) across both heaps.

Edge cases and pitfalls