Median of Data Stream Using Two Heaps in DSA C++

Mental Model

Keep the smaller half of numbers in one heap and the larger half in another heap to quickly find the middle value.

Analogy: Imagine two baskets: one holds smaller fruits, the other holds bigger fruits. The middle fruit is the one between these baskets.

MaxHeap (smaller half): 10 -> 5 -> 3 -> null
MinHeap (larger half): 12 -> 15 -> 20 -> null

Dry Run Walkthrough

Input: Stream: 5, 15, 1, 3

Goal: Find the median after each number is added

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

MaxHeap: 5 -> null
MinHeap: null

Why: Start by adding first number to max heap to hold smaller half

Step 2: Add 15 to min heap (larger half)

MaxHeap: 5 -> null
MinHeap: 15 -> null

Why: Add next number to min heap to hold larger half

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

MaxHeap: 5 -> 1 -> null
MinHeap: 15 -> null

Why: Add smaller number to max heap to keep smaller half

Step 4: No balancing needed

MaxHeap: 5 -> 1 -> null
MinHeap: 15 -> null

Why: maxHeap.size() = 2 == minHeap.size() + 1, which is allowed; median = maxHeap.top() = 5

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

MaxHeap: 5 -> 1 -> 3 -> null
MinHeap: 15 -> null

Why: 3 <= maxHeap.top() = 5, so add to smaller half

Step 6: Balance heaps to keep size difference ≤ 1

MaxHeap: 3 -> 1 -> null
MinHeap: 5 -> 15 -> null

Why: maxHeap.size() = 3 > minHeap.size() + 1 = 2, so move maxHeap.top() = 5 to minHeap and pop it from maxHeap

Result:

After each insertion, medians are: 5, 10, 5, 4
Final heaps:
MaxHeap: 3 -> 1 -> null
MinHeap: 5 -> 15 -> null

Annotated Code

DSA C++

#include <iostream>
#include <queue>
#include <vector>

using namespace std;

class MedianFinder {
private:
    priority_queue<int> maxHeap; // smaller half
    priority_queue<int, vector<int>, greater<int>> minHeap; // larger half

public:
    void addNum(int num) {
        if (maxHeap.empty() || num <= maxHeap.top()) {
            maxHeap.push(num); // add to smaller half
        } else {
            minHeap.push(num); // add to larger half
        }

        // Balance the heaps so that size difference is not more than 1
        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.push(maxHeap.top());
            maxHeap.pop();
        } else if (minHeap.size() > maxHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }

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

int main() {
    MedianFinder mf;
    vector<int> stream = {5, 15, 1, 3};

    for (int num : stream) {
        mf.addNum(num);
        cout << mf.findMedian() << endl;
    }

    return 0;
}

if (maxHeap.empty() || num <= maxHeap.top()) {

decide to add number to smaller half if empty or number is smaller or equal to max of smaller half

maxHeap.push(num);

add number to max heap (smaller half)

minHeap.push(num);

add number to min heap (larger half)

if (maxHeap.size() > minHeap.size() + 1) {

check if smaller half is too big

minHeap.push(maxHeap.top()); maxHeap.pop();

move top from smaller half to larger half to balance

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

check if larger half is bigger

maxHeap.push(minHeap.top()); minHeap.pop();

move top from larger half to smaller half to balance

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

if equal size, median is average of tops

return (maxHeap.top() + minHeap.top()) / 2.0;

calculate median as average

return maxHeap.top();

if sizes differ, median is top of bigger heap (smaller half)

Output

5

10

5

4

Complexity Analysis

Time: O(log n) per insertion because each addNum operation pushes/pops from heaps which take log n 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)), using two heaps is much faster for streaming data

Edge Cases

empty stream

findMedian should not be called or should handle gracefully

DSA C++

if (maxHeap.empty() || num <= maxHeap.top()) {

all numbers equal

heaps balance normally, median is that number

DSA C++

if (maxHeap.empty() || num <= maxHeap.top()) {

stream with one element

median is that element

DSA C++

if (maxHeap.empty() || num <= maxHeap.top()) {

When to Use This Pattern

When you need to find the median in a stream of numbers efficiently, use two heaps to keep track of smaller and larger halves for quick median calculation.

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 max heap and min heap based on their value relative to current median

Summary

It finds the median of a stream of numbers by using two heaps to store smaller and larger halves.

Use it when you need fast median queries on continuously incoming data.

The key insight is to keep two balanced heaps so the median is always at the top of one or average of both.

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 efficiently find the median by balancing smaller and larger halves of the data

B. To sort the entire data stream in ascending order

C. To store all elements in one heap for quick access

D. To remove duplicate elements from 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 how median is found

   Balancing these heaps allows quick access to the middle values, enabling efficient median calculation.

3. Final Answer:

   To efficiently find the median by balancing smaller and larger halves of the data -> Option A

4. Quick Check:

   Two heaps balance halves = median found fast [OK]

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

Common Mistakes:

• Thinking heaps sort entire data
• Using one heap only
• Confusing heaps with arrays

2. Which of the following is the default way to declare a max-heap in C++ using the priority_queue?

easy

A. priority_queue maxHeap;

B. priority_queue<int> maxHeap;

C. priority_queue maxHeap;

D. priority_queue maxHeap;

Solution

1. Step 1: Recall default priority_queue behavior

   By default, priority_queue in C++ is a max-heap using less<int>.

2. Step 2: Identify correct declaration

   Declaring priority_queue<int> maxHeap; creates a max-heap without extra parameters.

3. Final Answer:

   priority_queue<int> maxHeap; -> Option B

4. Quick Check:

   Default priority_queue = max-heap [OK]

Hint: Default priority_queue is max-heap [OK]

Common Mistakes:

• Using greater<int> for max-heap
• Omitting vector container
• Confusing max-heap with min-heap syntax

3. Consider the following C++ code snippet that inserts numbers into two heaps to find the median:

priority_queue<int> maxHeap;
priority_queue minHeap;

// Insert 10
maxHeap.push(10);
// Insert 20
minHeap.push(20);
// Insert 15
if (15 <= maxHeap.top()) maxHeap.push(15); else minHeap.push(15);

// Balance heaps
if (maxHeap.size() > minHeap.size() + 1) {
  minHeap.push(maxHeap.top());
  maxHeap.pop();
} else if (minHeap.size() > maxHeap.size()) {
  maxHeap.push(minHeap.top());
  minHeap.pop();
}

// Calculate median
if (maxHeap.size() == minHeap.size())
  cout << (maxHeap.top() + minHeap.top()) / 2.0 << endl;
else
  cout << maxHeap.top() << endl;

What is the output of this code?

medium

A. 15

B. 15.0

C. 12.5

D. 20

Solution

1. Step 1: Insert elements and balance heaps

   Insert 10 into maxHeap: maxHeap = [10]. Insert 20 into minHeap: minHeap = [20]. Insert 15: 15 > 10, so goes to minHeap: minHeap = [15, 20]. Balance heaps: minHeap size (2) > maxHeap size (1), so move minHeap top (15) to maxHeap. Now maxHeap = [15, 10], minHeap = [20].

2. Step 2: Calculate median

   maxHeap size = 2, minHeap size = 1, so median is maxHeap.top() = 15. Sizes unequal, cout << maxHeap.top() << endl; outputs 15.

3. Final Answer:

   15 -> Option A

4. Quick Check:

   Median after balancing = 15 [OK]

Hint: Balance heaps, median = top of larger heap or average [OK]

Common Mistakes:

• Not balancing heaps correctly
• Integer division instead of float
• Wrong heap insertion condition

4. The following code snippet is intended to maintain two heaps for median calculation. Identify the bug:

priority_queue<int> maxHeap;
priority_queue minHeap;

void addNumber(int num) {
  if (maxHeap.empty() || num <= maxHeap.top()) {
    maxHeap.push(num);
  } else {
    minHeap.push(num);
  }

  if (maxHeap.size() > minHeap.size()) {
    minHeap.push(maxHeap.top());
    maxHeap.pop();
  } else if (minHeap.size() > maxHeap.size() + 1) {
    maxHeap.push(minHeap.top());
    minHeap.pop();
  }
}

medium

A. The heaps are not declared correctly for median calculation

B. The comparison operator in if condition should be num >= maxHeap.top()

C. The balancing condition for maxHeap is incorrect; it should check if maxHeap.size() > minHeap.size() + 1

D. The function does not insert elements into any heap

Solution

1. Step 1: Analyze balancing conditions

   The code moves element from maxHeap to minHeap if maxHeap.size() > minHeap.size(), but it should only do so if maxHeap.size() > minHeap.size() + 1 to keep heaps balanced properly.

2. Step 2: Confirm correct balancing logic

   Proper balancing ensures size difference is at most 1. The current condition causes unnecessary moves, breaking balance.

3. Final Answer:

   The balancing condition for maxHeap is incorrect; it should check if maxHeap.size() > minHeap.size() + 1 -> Option C

4. Quick Check:

   Balance heaps only if size difference > 1 [OK]

Hint: Balance heaps only if size difference > 1 [OK]

Common Mistakes:

• Wrong balancing condition
• Swapping heap roles
• Incorrect comparison operators

5. You are given a stream of integers: [5, 15, 1, 3, 8, 7, 9, 10, 6, 11]. Using two heaps to find the median after each insertion, what is the median after inserting all elements?

hard

A. 8

B. 7

C. 6.5

D. 7.5

Solution

1. Step 1: Insert elements and balance heaps step-by-step

   After processing the entire stream with proper insertion (to maxHeap if <= max.top or empty, else minHeap) and balancing (rebalance if max.size > min.size + 1 or min.size > max.size), final state: maxHeap size=5 (top=7), minHeap size=5 (top=8).

2. Step 2: Calculate final median

   After all insertions, maxHeap size = 5, minHeap size = 5, median = average of tops: (7 + 8) / 2 = 7.5

3. Final Answer:

   7.5 -> Option D

4. Quick Check:

   Median after all insertions = 7.5 [OK]

Hint: Median = average of two heap tops if equal size [OK]

Common Mistakes:

• Not balancing heaps after each insertion
• Calculating median from wrong heap
• Forgetting to average when heaps equal size