Topics › Heaps › Hard Problems

Find Median from Data Stream

Hard TUF ID 566 A2ZSDEBlind 75

LeetCode ↗ Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

Design a MedianFinder that receives integers one at a time and can report the median of everything received so far:

• addNum(num) adds an integer from the stream.
• findMedian() returns the median as a double. With an odd count it is the middle value in sorted order. With an even count it is the average of the two middle values.

Constraints: -10^5 ≤ num ≤ 10^5, at least one number is added before findMedian is called, and there are at most 5 · 10^4 calls. The code here also handles the full int range and throws IllegalStateException on 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 implement one 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 store everything, and when asked, sort a copy and read the middle. Adding is trivial. Every query redoes a full sort, even when only one number has arrived since the last one.

Algorithm.

1. addNum: append to a list.
2. findMedian: copy into an array and sort. If the size n is odd, return a[n/2]. Otherwise return (a[n/2 - 1] + a[n/2]) / 2.0, adding in long.

static class SortEachQuery 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("no numbers yet");
        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;
    }
}

Complexity. addNum takes O(1). findMedian takes O(n log n). Space O(n).

If queries are frequent, for example one after every add, the total is O(n^2 log n).

Approach 2: Better, keep the list sorted

Intuition. Instead of sorting at query time, keep the list sorted at all times. Each new number goes into its correct position, found by binary search. The median is then just an index lookup.

Algorithm.

1. addNum: binary search for the first index whose value is greater than num, then insert num there.
2. findMedian: read the middle element, or the two middle elements, directly.

static class SortedInsert implements MedianFinder {
    private final List<Integer> sorted = new ArrayList<>();

    public void addNum(int num) {
        int lo = 0, hi = sorted.size();                     // first index whose value is > num
        while (lo < hi) {
            int mid = (lo + hi) >>> 1;
            if (sorted.get(mid) <= num) lo = mid + 1; else hi = mid;
        }
        sorted.add(lo, num);
    }

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

Complexity. addNum takes O(n). The binary search is only O(log n), but ArrayList.add(index, value) shifts every element after the insertion point. findMedian takes O(1). Space O(n).

This is fine for a few thousand numbers. The shifting is the bottleneck, and it exists because we maintain a total order. The median only needs to know which numbers are on the left of the middle and which are on the right, not the order inside each side.

Approach 3: Optimal, two heaps

Intuition. Split the numbers into two halves:

• low, a max-heap holding the smaller half. Its top is the largest of the small numbers.
• high, a min-heap holding the larger half. Its top is the smallest of the large numbers.

If every number in low is ≤ every number in high, and the halves have equal sizes (or low has one extra), then the median sits right at the boundary. It is low.peek() when the count is odd, and the average of low.peek() and high.peek() when it is even. Neither half is ever sorted internally. The heaps only expose their extreme elements, which is all the median needs.

Adding a number takes three moves:

1. Push num into low.
2. Move low’s maximum into high. Whatever num was, the largest element of the lower side now crosses over, so the ordering between the halves holds again: everything still in low was ≤ the element that moved and ≤ everything already in high.
3. If high is now larger than low, move high’s minimum back into low. This restores the size rule without breaking the ordering.

Algorithm.

1. addNum(num): low.add(num), then high.add(low.poll()), then if high.size() > low.size(), low.add(high.poll()).
2. findMedian(): if low is bigger, return low.peek(). Otherwise return (low.peek() + high.peek()) / 2.0, adding in long.

static class TwoHeaps implements MedianFinder {
    private final PriorityQueue<Integer> low = new PriorityQueue<>(Collections.reverseOrder());   // max-heap
    private final PriorityQueue<Integer> high = new PriorityQueue<>();                            // min-heap

    public void addNum(int num) {
        low.add(num);
        high.add(low.poll());                               // largest of the lower half crosses over
        if (high.size() > low.size()) low.add(high.poll()); // keep low the same size or one bigger
    }

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

Dry run on example 2:

Complexity. addNum takes O(log n): at most five heap operations. findMedian takes O(1). Space O(n).

Follow-ups often asked. If every number lies in [0, 100], keep a count array of size 101. Adding is O(1) and the median is a walk over 101 counters. If 99% of numbers lie in that range, keep the count array and add two counters for “below 0” and “above 100”. The median almost surely falls inside the array, and the rare outliers only need counting.

Edge cases and pitfalls

• Integer division. (a + b) / 2 on ints drops the .5. Divide by 2.0.
• Overflow in the average. Integer.MAX_VALUE + Integer.MAX_VALUE wraps to -2. Cast to long before adding. The tests add MAX_VALUE twice and expect 2147483647.0, and add MIN_VALUE and MAX_VALUE and expect -0.5.
• Comparator overflow. A max-heap written as (a, b) -> b - a breaks for extreme values, because b - a overflows and flips the sign. Use Collections.reverseOrder() or Integer.compare(b, a).
• Which heap holds the extra element. Pick one convention (here low) and use it consistently in both addNum and findMedian. Mixing conventions returns the wrong middle for odd counts.
• Inserting straight into the “right” heap by comparing with low.peek() also works, but it needs care when low is empty. The three-move version never compares anything, which leaves fewer branches to get wrong.
• Empty stream. LeetCode never queries an empty finder. The code throws an exception rather than returning a misleading 0.0.
• Duplicates need no special handling. Equal values may end up on either side of the boundary.

Related problems

• Kth largest element in a stream of running integers keeps one heap over a stream instead of two.
• Find Median in a Stream is the SDE-sheet listing of the same task.
• Median of 2 sorted arrays finds a median across two sorted arrays, using binary search instead of heaps.
• Matrix Median finds the median of a row-sorted matrix by binary searching on the answer.
• Implement Max Heap shows how the low heap works internally.