Median of Data Stream Using Two Heaps in DSA Java

Mental Model

Keep the smaller half of numbers in a max heap and the larger half in a min heap to quickly find the middle value.

Analogy: Imagine two balanced baskets: one holds smaller fruits, the other holds bigger fruits. The median is the fruit right between the two baskets.

MaxHeap (smaller half): [5, 3, 1]
MinHeap (larger half): [6, 8, 9]
Median is between the top of both heaps.

Dry Run Walkthrough

Input: Stream: 5, 3, 8, 1, 6

Goal: Find the median after each number is added

Step 1: Add 5 to max heap (smaller half)

MaxHeap: [5]
MinHeap: []

Why: Start by adding the first number to max heap

Step 2: Add 3 to max heap, then balance by moving top to min heap

MaxHeap: [3]
MinHeap: [5]

Why: 3 <= maxHeap.peek() (5), add to max heap; balance since maxHeap.size() - minHeap.size() > 1, move maxHeap top (largest) to min heap

Step 3: Add 8 to min heap, then balance by moving top to max heap

MaxHeap: [5, 3]
MinHeap: [8]

Why: 8 > maxHeap.peek() (3), add to min heap; balance since maxHeap.size() < minHeap.size(), move minHeap top (smallest) to max heap

Step 4: Add 1 to max heap, then balance by moving top to min heap

MaxHeap: [3, 1]
MinHeap: [5, 8]

Why: 1 <= maxHeap.peek() (5), add to max heap; balance since maxHeap.size() - minHeap.size() > 1, move maxHeap top (largest, 5) to min heap

Step 5: Add 6 to min heap, then balance by moving top to max heap

MaxHeap: [5, 3, 1]
MinHeap: [6, 8]

Why: 6 > maxHeap.peek() (3), add to min heap; balance since maxHeap.size() < minHeap.size(), move minHeap top (smallest, 5) to max heap

Result:

After each insertion, median is:
5.0
4.0
5.0
4.0
5.0

Annotated Code

DSA Java

import java.util.PriorityQueue;

public class MedianFinder {
    private PriorityQueue<Integer> maxHeap; // smaller half
    private PriorityQueue<Integer> minHeap; // larger half

    public MedianFinder() {
        maxHeap = new PriorityQueue<>((a, b) -> b - a); // max heap
        minHeap = new PriorityQueue<>(); // min heap
    }

    public void addNum(int num) {
        if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
            maxHeap.offer(num); // add to smaller half
        } else {
            minHeap.offer(num); // add to larger half
        }

        // Balance heaps so that maxHeap size >= minHeap size and difference <= 1
        if (maxHeap.size() < minHeap.size()) {
            maxHeap.offer(minHeap.poll()); // move smallest from minHeap to maxHeap
        } else if (maxHeap.size() - minHeap.size() > 1) {
            minHeap.offer(maxHeap.poll()); // move largest from maxHeap to minHeap
        }
    }

    public double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            return (maxHeap.peek() + minHeap.peek()) / 2.0; // average of two middles
        } else {
            return maxHeap.peek(); // maxHeap has one extra element
        }
    }

    public static void main(String[] args) {
        MedianFinder mf = new MedianFinder();
        int[] stream = {5, 3, 8, 1, 6};
        for (int num : stream) {
            mf.addNum(num);
            System.out.println(mf.findMedian());
        }
    }
}

if (maxHeap.isEmpty() || num <= maxHeap.peek()) {

decide which heap to add the new number to keep smaller half in maxHeap

maxHeap.offer(num);

add number to maxHeap if smaller or equal to maxHeap top

minHeap.offer(num);

add number to minHeap if greater than maxHeap top

if (maxHeap.size() < minHeap.size()) {

balance heaps if minHeap has more elements

maxHeap.offer(minHeap.poll());

move smallest from minHeap to maxHeap to balance

else if (maxHeap.size() - minHeap.size() > 1) {

balance heaps if maxHeap has more than one extra element

minHeap.offer(maxHeap.poll());

move largest from maxHeap to minHeap to balance

if (maxHeap.size() == minHeap.size()) {

if heaps equal size, median is average of tops

return maxHeap.peek();

if maxHeap has one extra element, median is its top

Output

5.0

4.0

5.0

4.0

5.0

Complexity Analysis

Time: O(log n) per insertion because adding to heaps and balancing takes logarithmic time

Space: O(n) because all numbers are stored in the two heaps

vs Alternative: Compared to sorting all numbers each time (O(n log n)), this approach is faster for continuous median queries

Edge Cases

Empty stream

No median can be found; code does not call findMedian before adding numbers

DSA Java

if (maxHeap.isEmpty() || num <= maxHeap.peek()) {

All numbers equal

Heaps balance normally; median is that number

DSA Java

if (maxHeap.isEmpty() || num <= maxHeap.peek()) {

Stream with one element

Median is that single element

DSA Java

if (maxHeap.size() == minHeap.size()) {

When to Use This Pattern

When asked to find the median of a growing list of numbers efficiently, use two heaps to keep track of smaller and larger halves for quick median retrieval.

Common Mistakes

Mistake: Not balancing the heaps after insertion, causing incorrect median calculation
Fix: Always balance heaps so their sizes differ by at most one after each insertion

Mistake: Adding all numbers to one heap only
Fix: Distribute numbers between maxHeap and minHeap based on their value relative to current median

Summary

Maintains two heaps to efficiently find the median of a data stream.

Use when you need the median after each new number in a stream without sorting all numbers repeatedly.

Keep smaller half in max heap and larger half in min heap, balancing sizes to find median quickly.

Practice

1. What is the main purpose of using two heaps (a max heap and a min heap) in the Median of Data Stream problem?

easy

A. To store all numbers in a single data structure

B. To sort all numbers in ascending order

C. To efficiently keep track of the lower half and upper half of the numbers for median calculation

D. To find the maximum number in the data stream

Solution

1. Step 1: Understand the role of two heaps

   The max heap stores the smaller half of the numbers, and the min heap stores the larger half.

2. Step 2: Explain median calculation

   By balancing these heaps, the median can be found quickly from the top elements of the heaps.

3. Final Answer:

   To efficiently keep track of the lower half and upper half of the numbers for median calculation -> Option C

4. Quick Check:

   Two heaps split data for median [OK]

Hint: Two heaps split data into halves for quick median [OK]

Common Mistakes:

• Thinking heaps sort all numbers completely
• Using only one heap for median
• Confusing max heap and min heap roles

2. Which of the following Java code snippets correctly initializes a max heap using PriorityQueue for the Median of Data Stream problem?

easy

A. PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

B. PriorityQueue<Integer> maxHeap = new PriorityQueue<>();

C. PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.naturalOrder());

D. PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.naturalOrder());

Solution

1. Step 1: Identify max heap initialization

   Java's PriorityQueue is a min heap by default; to make max heap, use Collections.reverseOrder().

2. Step 2: Check options

   PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder()); uses Collections.reverseOrder(), which correctly creates a max heap.

3. Final Answer:

   PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder()); -> Option A

4. Quick Check:

   Max heap needs reverse order comparator [OK]

Hint: Use Collections.reverseOrder() for max heap in Java [OK]

Common Mistakes:

• Using default PriorityQueue without comparator
• Using naturalOrder() which is min heap
• Using Collections.naturalOrder() which doesn't exist

3. Consider the following Java code snippet for adding numbers to two heaps and finding the median:

MedianFinder mf = new MedianFinder();
mf.addNum(1);
mf.addNum(5);
mf.addNum(2);
System.out.println(mf.findMedian());

Assuming the MedianFinder class uses two heaps correctly, what will be the printed output?

medium

A. 2.5

B. 2

C. 3

D. 1

Solution

1. Step 1: Add numbers and balance heaps

   Numbers added: 1, 5, 2. Sorted order: 1, 2, 5. Median is middle number 2.

2. Step 2: Calculate median

   Since odd count, median is the root of the heap with one extra element, which is 2.

3. Final Answer:

   2 -> Option B

4. Quick Check:

   Median of [1,2,5] is 2 [OK]

Hint: Median of odd count is middle element [OK]

Common Mistakes:

• Calculating average instead of middle for odd count
• Ignoring heap balancing
• Confusing median with mean

4. The following Java method is part of a MedianFinder class using two heaps. It is supposed to add a number and rebalance heaps. Identify the bug:

public void addNum(int num) {
    if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
        maxHeap.add(num);
    } else {
        minHeap.add(num);
    }
    if (maxHeap.size() > minHeap.size() + 1) {
        minHeap.add(maxHeap.poll());
    } else if (minHeap.size() > maxHeap.size()) {
        maxHeap.add(minHeap.poll());
    }
}

medium

A. Should check if maxHeap is empty before peek() to avoid NullPointerException

B. The addNum method should not add to minHeap when num equals maxHeap.peek()

C. The rebalance condition for minHeap size is incorrect; should be minHeap.size() > maxHeap.size() + 1

D. No bug, method works correctly

Solution

1. Step 1: Analyze addNum logic

   It adds to maxHeap if empty or num <= maxHeap.peek(), else to minHeap.

2. Step 2: Check rebalance conditions

   Rebalances if maxHeap is bigger by more than 1 or minHeap bigger than maxHeap.

3. Final Answer:

   No bug, method works correctly -> Option D

4. Quick Check:

   Correct heap balancing logic [OK]

Hint: Check heap sizes differ by at most one [OK]

Common Mistakes:

• Forgetting to check maxHeap empty before peek
• Incorrect rebalance size conditions
• Misplacing elements between heaps

5. You want to extend the MedianFinder class to handle a large data stream efficiently. Which approach improves performance when finding the median after every insertion?

hard

A. Use two balanced heaps to keep halves and rebalance after each insertion

B. Store all numbers in a list and sort it after each insertion

C. Use a single max heap to store all numbers

D. Use a hash map to count frequencies and calculate median from counts

Solution

1. Step 1: Evaluate approaches for median calculation

   Sorting after each insertion (Store all numbers in a list and sort it after each insertion) is costly (O(n log n)). Single heap or hash map do not support quick median.

2. Step 2: Explain two heaps efficiency

   Two heaps keep halves balanced, allowing O(log n) insertion and O(1) median retrieval.

3. Final Answer:

   Use two balanced heaps to keep halves and rebalance after each insertion -> Option A

4. Quick Check:

   Two heaps optimize median in data stream [OK]

Hint: Two heaps give fast median updates in streams [OK]

Common Mistakes:

• Sorting entire list each time
• Using single heap only
• Ignoring insertion cost