Skip to content
AI360Xpert
Back to Heap / Priority Queue
Hard

Find Median from Data Stream

The median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value, and the median is the mean of the two middle values. Implement the `MedianFinder` class.

Examples

Input:["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"] [[], [1], [2], [], [3], []]
Output:[null, null, null, 1.5, null, 2.0]
MedianFinder medianFinder = new MedianFinder(); medianFinder.addNum(1); // arr = [1] medianFinder.addNum(2); // arr = [1, 2] medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2) medianFinder.addNum(3); // arr[1, 2, 3] medianFinder.findMedian(); // return 2.0

Constraints

  • -10^5 <= num <= 10^5
  • There will be at least one element in the data structure before calling findMedian.
  • At most 5 * 10^4 calls will be made to addNum and findMedian.

Two Heaps

Approach

Maintain two heaps: a Max Heap (`small`) to store the smaller half of the numbers, and a Min Heap (`large`) to store the larger half. Maintain two invariants: 1) Every element in `small` <= every element in `large`. 2) The size difference between `small` and `large` is at most 1 (we can choose `small` to have the extra element if the total count is odd). For `addNum`, push to `small`, then move the max of `small` to `large` to maintain invariant 1. If `large` becomes larger than `small`, move the min of `large` back to `small` to maintain invariant 2. For `findMedian`, if sizes are unequal, return the root of `small`. If sizes are equal, return the average of the roots of both heaps.

Complexity Analysis

Time Complexity
O(log N) add, O(1) findMedian
Space Complexity
O(N)

This optimally maintains the median over a stream of incoming data.

Solution.java
class MedianFinder {    private PriorityQueue<Integer> small; // Max Heap for smaller half    private PriorityQueue<Integer> large; // Min Heap for larger half
    public MedianFinder() {        small = new PriorityQueue<>(Collections.reverseOrder());        large = new PriorityQueue<>();    }        public void addNum(int num) {        // Add to small heap first        small.offer(num);                // Ensure every element in small is <= every element in large        if (!small.isEmpty() && !large.isEmpty() && small.peek() > large.peek()) {            large.offer(small.poll());        }                // Balance sizes: small heap can have at most 1 more element than large heap        if (small.size() > large.size() + 1) {            large.offer(small.poll());        } else if (large.size() > small.size()) {            small.offer(large.poll());        }    }        public double findMedian() {        if (small.size() == large.size()) {            return (small.peek() + large.peek()) / 2.0;        } else {            return small.peek();        }    }}