Median of Data Stream Using Two Heaps in DSA Java - Build from Scratch

Median of Data Stream Using Two Heaps

📖 Scenario: Imagine you are building a system that receives numbers one by one, like scores from a game played live. You want to quickly find the middle score (median) at any time without sorting all scores every time a new number arrives.

🎯 Goal: Build a Java program that uses two heaps (priority queues) to keep track of the median of a stream of numbers efficiently.

📋 What You'll Learn

Create two heaps: a max heap for the lower half of numbers and a min heap for the upper half.

Add numbers one by one to the correct heap to keep them balanced.

Calculate the median based on the sizes of the two heaps.

Print the median after all numbers are added.

💡 Why This Matters

🌍 Real World

This technique is used in live data analysis, like tracking median latency in networks or median prices in stock markets.

💼 Career

Understanding heaps and median calculation is useful for software engineers working on real-time analytics, financial software, or performance monitoring tools.

1

Create two heaps for storing numbers

Create a PriorityQueue maxHeap as a max heap using a comparator, and a PriorityQueue minHeap as a min heap.

DSA Java

// Create maxHeap and minHeap
// Your code here

Hint

Use Collections.reverseOrder() to create a max heap in Java.

2

Add a helper method to add numbers to heaps

Add a method public void addNum(int num) that adds num to maxHeap if it is empty or num is less than or equal to the top of maxHeap, otherwise add it to minHeap. Then balance the heaps so their sizes differ by at most 1.

DSA Java

import java.util.PriorityQueue;
import java.util.Collections;

public class MedianFinder {
    PriorityQueue maxHeap = new PriorityQueue(Collections.reverseOrder());
    PriorityQueue minHeap = new PriorityQueue();

    // Add addNum method below
    // Your code here
}

Hint

Check if maxHeap is empty or num is smaller than or equal to maxHeap.peek() to decide where to add.

Balance heaps by moving elements between them if size difference is more than 1.

3

Add a method to find the median

Add a method public double findMedian() that returns the median: if heaps are equal size, return average of maxHeap.peek() and minHeap.peek(), else return maxHeap.peek().

DSA Java

import java.util.PriorityQueue;
import java.util.Collections;

public class MedianFinder {
    PriorityQueue maxHeap = new PriorityQueue(Collections.reverseOrder());
    PriorityQueue minHeap = new PriorityQueue();

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

    // Add findMedian method below
    // Your code here
}

Hint

If both heaps have the same size, median is average of their top elements. Otherwise, median is top of maxHeap.

4

Add numbers and print the median

In the main method, create a MedianFinder object, add these numbers in order: 5, 15, 1, 3 using addNum, then print the median using findMedian().

DSA Java

import java.util.PriorityQueue;
import java.util.Collections;

public class MedianFinder {
    PriorityQueue maxHeap = new PriorityQueue(Collections.reverseOrder());
    PriorityQueue minHeap = new PriorityQueue();

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

    public double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            return (maxHeap.peek() + minHeap.peek()) / 2.0;
        } else {
            return maxHeap.peek();
        }
    }

    public static void main(String[] args) {
        // Create MedianFinder and add numbers 5, 15, 1, 3
        // Print the median
        // Your code here
    }
}

Hint

After adding numbers 5, 15, 1, 3, the median should be 4.0.

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