Find Median from Data Stream

Description

295. 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.

• For example, for arr = [2,3,4], the median is 3.
• For example, for arr = [2,3], the median is (2 + 3) / 2 = 2.5.

Implement the MedianFinder class:

• MedianFinder() initializes the MedianFinder object.
• void addNum(int num) adds the integer num from the data stream to the data structure.
• double findMedian() returns the median of all elements so far. Answers within $10^{-5}$ of the actual answer will be accepted.

Example 1:

Input
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
Output
[null, null, null, 1.5, null, 2.0]
 
Explanation
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\times 10^4$ calls will be made to addNum and findMedian.

---

Brute Force Approach: Insertion Sort / List Sorting

Intuition

Maintain a sorted dynamic array. Every time a new number arrives, insert it into its sorted position using binary search. To find the median, retrieve the element(s) at index n / 2.

import java.util.*;
 
class MedianFinder {
    private List<Integer> list;
 
    public MedianFinder() {
        list = new ArrayList<>();
    }
    
    public void addNum(int num) {
        int index = Collections.binarySearch(list, num);
        if (index < 0) {
            index = -(index + 1);
        }
        list.add(index, num); // Insertion takes O(N) shift time
    }
    
    public double findMedian() {
        int n = list.size();
        if (n % 2 == 1) {
            return list.get(n / 2);
        } else {
            return (list.get(n / 2 - 1) + list.get(n / 2)) / 2.0;
        }
    }
}
 

Complexity

• Time Complexity:
  • addNum: $O(N)$ — Binary search takes $O(\log N)$, but array element shifting takes $O(N)$.
  • findMedian: $O(1)$ — Direct index access.
• Space Complexity: $O(N)$ — Stores all elements in an array list.

---

Most Optimized Solution: Two Heaps (Max-Heap + Min-Heap)

Intuition

Divide all numbers into two equal-sized halves:

1. small (Max-Heap): Stores the smaller half of numbers (the largest value in this half is at the top).
2. large (Min-Heap): Stores the larger half of numbers (the smallest value in this half is at the top).

By balancing both heaps:

• Every element in small is $\le$ every element in large.
• The median is always accessible in $O(1)$ time at the top of the heaps.

import java.util.PriorityQueue;
import java.util.Collections;
 
class MedianFinder {
    private PriorityQueue<Integer> small; // Max-Heap for the smaller half
    private PriorityQueue<Integer> large; // Min-Heap for the larger half
 
    public MedianFinder() {
        small = new PriorityQueue<>(Collections.reverseOrder());
        large = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        // Step 1: Add to small heap
        small.add(num);
 
        // Step 2: Ensure max(small) <= min(large)
        large.add(small.poll());
 
        // Step 3: Keep small size >= large size (small can have at most 1 extra element)
        if (large.size() > small.size()) {
            small.add(large.poll());
        }
    }
    
    public double findMedian() {
        if (small.size() > large.size()) {
            return small.peek();
        } else {
            return (small.peek() + large.peek()) / 2.0;
        }
    }
}
 

Complexity

• Time Complexity:

• addNum: $O(\log N)$ — Heap push and pop operations.

• findMedian: $O(1)$ — Peek operations at the top of the heaps.

• Space Complexity: $O(N)$ — Stores all numbers across both heaps.

---

Follow-Up Solutions

1. All numbers in range [0, 100]

Use a frequency array of size 101 (count[101]) and a total element counter:

• addNum: Increment count[num] in $O(1)$ time.
• findMedian: Iterate through count[101] to find the 50th percentile index in $O(101)=O(1)$ time and $O(1)$ space.

2. 99% of numbers in range [0, 100]

Use a frequency array of size 101 for values in [0, 100] and two counters/structures for values strictly < 0 and > 100. Since 99% of elements lie within [0, 100], the median will almost certainly fall within [0, 100], maintaining $O(1)$ average time.

---

Easy Memory Rule

“Max-Heap for smaller half, Min-Heap for larger half. Keep small size $\ge$ large size.”

1. What is a Heap?

A Heap is a tree-based data structure optimized for fast access to either the highest or lowest value in a dataset.

• Max-Heap: The parent node is always greater than or equal to its children. The largest element sits at the root (top).
• Min-Heap: The parent node is always smaller than or equal to its children. The smallest element sits at the root (top).

Key Operations:

• peek() / top() $\to$ $O(1)$ (retrieve max/min without removing)
• add() / poll() $\to$ $O(\log N)$ (insert or remove the top element and re-balance tree)

---

2. The Two Heaps Used in Find Median from Data Stream

To find the median in $O(1)$ time, we divide all numbers into two equal-sized halves:

1. small (Max-Heap): Stores the smaller half of numbers.

   • Purpose: Holds the left half of the sorted sequence. Calling small.peek() gives the largest number of this lower half.

2. large (Min-Heap): Stores the larger half of numbers.

   • Purpose: Holds the right half of the sorted sequence. Calling large.peek() gives the smallest number of this upper half.

Core Invariant:
Every element in small must be $\le$ every element in large.

---

3. Step-by-Step Method Breakdown

import java.util.PriorityQueue;
import java.util.Collections;
 
class MedianFinder {
    private PriorityQueue<Integer> small; // Max-Heap
    private PriorityQueue<Integer> large; // Min-Heap
 
    public MedianFinder() {
        small = new PriorityQueue<>(Collections.reverseOrder());
        large = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        small.add(num);
        large.add(small.poll());
 
        if (large.size() > small.size()) {
            small.add(large.poll());
        }
    }
    
    public double findMedian() {
        if (small.size() > large.size()) {
            return small.peek();
        } else {
            return (small.peek() + large.peek()) / 2.0;
        }
    }
}
 

Method 1: MedianFinder() (Constructor)

• Goal: Initialize both heaps.
• small uses Collections.reverseOrder() so that standard Java PriorityQueue acts as a Max-Heap.
• large uses default behavior, acting as a Min-Heap.

Method 2: addNum(int num)

• Goal: Add num and keep both heaps sorted relative to each other and balanced in size.
• Step 1 (small.add(num)): Push num into the small Max-Heap.
• Step 2 (large.add(small.poll())): Remove the largest value from small and push it to large. This guarantees that all elements in small are $\le$ all elements in large.
• Step 3 (if (large.size() > small.size()) small.add(large.poll())): If large has more elements than small, move its smallest element back to small. This ensures small carries the odd extra element whenever total count is odd (small.size() is either equal to or $1$ greater than large.size()).

Method 3: findMedian()

• Goal: Return the median in $O(1)$ time.

• Odd total elements (small.size() > large.size()):

• The median is the single extra middle element stored at the top of small (small.peek()).

• Even total elements (small.size() == large.size()):

• The median is the average of the two middle elements: the largest of the lower half (small.peek()) and the smallest of the upper half (large.peek()).

Comparison

1. Odd Count: [1, 2, 5, 8, 9] (5 elements)

$\begin{array}{ccc}\text{Lower Half}& \text{MIDDLE}& \text{Upper Half}\\ [1,2]& 5& [8,9]\end{array}$

• There is a physical middle element: 5.
• Median = 5

2. Even Count: [1, 2, 5, 8] (4 elements)

$\begin{array}{ccc}\text{Lower Half}& \text{GAP}& \text{Upper Half}\\ [1,2]& \uparrow & [5,8]\end{array}$

• The two middle elements are 2 and 5.
• Neither 2 nor 5 alone is the middle line.
  • If you pick 2, only 1 element is smaller (1), but 2 elements are bigger (5, 8). That’s unbalanced!
  • If you pick 5, 2 elements are smaller (1, 2), but only 1 element is bigger (8). Also unbalanced!
• To find the exact midpoint of the gap between 2 and 5, you calculate their average:
  $\text{Median}=\frac{2+5}{2}=3.5$

The Role of the Average

The average ($\frac{2+5}{2}=3.5$) creates a synthetic “center point”:

• 2 elements are smaller than 3.5 (1, 2).
• 2 elements are larger than 3.5 (5, 8).

So for an even number of elements, “average of the two middle elements” is just the mathematical rule used to produce the single dividing value.