Median of Data Stream Using Two Heaps in DSA Rust

Mental Model

We keep two heaps to split numbers into smaller half and larger half, so the middle value is easy to find.

Analogy: Imagine a seesaw balanced in the middle. The smaller numbers sit on the left side, the bigger numbers on the right side. The middle is where the seesaw balances.

MaxHeap (smaller half)       MinHeap (larger half)
      [10, 8, 5]                  [12, 15, 20]
          ↑                           ↑
       max at top               min at top

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]
MinHeap: []
Median: 5

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

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

MaxHeap: [5]
MinHeap: [15]
Median: (5 + 15) / 2 = 10

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

Step 3: Add 1 to max heap (smaller half), then rebalance heaps

MaxHeap: [5, 1]
MinHeap: [15]
Median: 5

Why: Add smaller number to max heap and keep heaps balanced in size

Step 4: Add 3 to max heap, then move max from max heap to min heap to rebalance

MaxHeap: [3, 1]
MinHeap: [5, 15]
Median: (3 + 5) / 2 = 4

Why: Keep heaps balanced by moving largest from smaller half to larger half

Result:

MaxHeap: [3, 1]
MinHeap: [5, 15]
Median: 4

Annotated Code

DSA Rust

use std::collections::BinaryHeap;
use std::cmp::Reverse;

struct MedianFinder {
    max_heap: BinaryHeap<i32>,       // smaller half
    min_heap: BinaryHeap<Reverse<i32>>, // larger half
}

impl MedianFinder {
    fn new() -> Self {
        MedianFinder {
            max_heap: BinaryHeap::new(),
            min_heap: BinaryHeap::new(),
        }
    }

    fn add_num(&mut self, num: i32) {
        if self.max_heap.peek().map_or(true, |&x| num <= x) {
            self.max_heap.push(num); // add to smaller half
        } else {
            self.min_heap.push(Reverse(num)); // add to larger half
        }

        // Rebalance heaps if size difference > 1
        if self.max_heap.len() > self.min_heap.len() + 1 {
            if let Some(max_top) = self.max_heap.pop() {
                self.min_heap.push(Reverse(max_top));
            }
        } else if self.min_heap.len() > self.max_heap.len() {
            if let Some(Reverse(min_top)) = self.min_heap.pop() {
                self.max_heap.push(min_top);
            }
        }
    }

    fn find_median(&self) -> f64 {
        if self.max_heap.len() > self.min_heap.len() {
            *self.max_heap.peek().unwrap() as f64
        } else {
            (*self.max_heap.peek().unwrap() as f64 + self.min_heap.peek().unwrap().0 as f64) / 2.0
        }
    }
}

fn main() {
    let mut mf = MedianFinder::new();
    let stream = vec![5, 15, 1, 3];

    for num in stream {
        mf.add_num(num);
        println!("Median after adding {}: {}", num, mf.find_median());
    }
}

if self.max_heap.peek().map_or(true, |&x| num <= x) {

decide which heap to add the new number to keep halves separated

self.max_heap.push(num);

add number to max heap (smaller half)

self.min_heap.push(Reverse(num));

add number to min heap (larger half)

if self.max_heap.len() > self.min_heap.len() + 1 {

check if max heap is too big and rebalance

if let Some(max_top) = self.max_heap.pop() {

move top from max heap to min heap to balance sizes

if self.min_heap.len() > self.max_heap.len() {

check if min heap is bigger and rebalance

if let Some(Reverse(min_top)) = self.min_heap.pop() {

move top from min heap to max heap to balance sizes

if self.max_heap.len() > self.min_heap.len() {

if max heap bigger, median is top of max heap

(*self.max_heap.peek().unwrap() as f64 + self.min_heap.peek().unwrap().0 as f64) / 2.0

if equal size, median is average of tops

Output

Median after adding 5: 5

Median after adding 15: 10

Median after adding 1: 5

Median after adding 3: 4

Complexity Analysis

Time: O(log n) per insertion because heap push/pop operations take logarithmic time

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

vs Alternative: Compared to sorting the entire stream each time (O(n log n)), using two heaps is much faster for continuous median queries

Edge Cases

Empty stream

Median cannot be found; code assumes at least one number before calling find_median

DSA Rust

No explicit guard; user must add numbers before calling find_median

All numbers equal

Heaps balance normally; median is that number

DSA Rust

if self.max_heap.peek().map_or(true, |&x| num <= x) {

Stream with negative and positive numbers

Heaps handle negative numbers correctly; median computed normally

DSA Rust

if self.max_heap.peek().map_or(true, |&x| num <= x) {

When to Use This Pattern

When asked to find median continuously from a stream of numbers, use two heaps to keep track of smaller and larger halves for efficient median retrieval.

Common Mistakes

Mistake: Not rebalancing heaps after insertion, causing size difference > 1
Fix: Always rebalance heaps after adding a number to keep their sizes balanced

Mistake: Adding all numbers to one heap without condition
Fix: Compare new number with max heap top to decide correct heap for insertion

Summary

Keeps track of median in a data stream using two heaps to split smaller and larger halves.

Use when you need to find median efficiently after each new number in a stream.

The key is balancing two heaps so the median is always at the top of one or average of tops.

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 sort the entire data stream before finding the median

B. To store all elements in one heap and ignore the other

C. To efficiently find the median by balancing smaller and larger halves of the data

D. To use heaps only for insertion speed without median calculation

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: How this helps find the median

   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 C

4. Quick Check:

   Two heaps balance halves = median found fast [OK]

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

Common Mistakes:

• Thinking heaps sort all data fully
• Using only one heap for median
• Ignoring balancing heaps after insertion

2. Which Rust data structure is best suited to implement a max-heap for the Median of Data Stream problem?

easy

A. BinaryHeap (max-heap by default)

B. HashMap storing counts of elements

C. Vec with manual sorting after each insertion

D. LinkedList for fast insertion and removal

Solution

1. Step 1: Identify Rust's heap implementation

   Rust's standard library provides BinaryHeap which is a max-heap by default.

2. Step 2: Using Reverse wrapper for min-heap

   To create a min-heap, BinaryHeap with Reverse wrapper is used; for max-heap, BinaryHeap alone suffices.

3. Final Answer:

   BinaryHeap (max-heap by default) -> Option A

4. Quick Check:

   BinaryHeap = max-heap in Rust [OK]

Hint: Rust BinaryHeap is max-heap by default [OK]

Common Mistakes:

• Using Vec and sorting every time (slow)
• Confusing HashMap with heap
• Using LinkedList which is not a heap

3. What will be the median after inserting the numbers 5, 15, 1, 3 in order using two heaps?

medium

A. 6

B. 5

C. 3

D. 4

Solution

1. Step 1: Insert 5 and find median

   Max-heap: [5], Min-heap: [], median = 5

2. Step 2: Insert 15 and balance heaps

   Max-heap: [5], Min-heap: [15], median = (5+15)/2 = 10

3. Step 3: Insert 1 and rebalance

   Max-heap: [5,1], Min-heap: [15], median = 5

4. Step 4: Insert 3 and rebalance

   Max-heap: [3,1], Min-heap: [5,15], median = (3+5)/2 = 4

5. Final Answer:

   4 -> Option D

6. Quick Check:

   Median after 4 inserts = 4 [OK]

Hint: Balance heaps after each insert, median is middle value(s) [OK]

Common Mistakes:

• Not balancing heaps after insertion
• Taking median as last inserted number
• Incorrect average calculation for even count

4. Given this Rust snippet for adding a number to heaps, what is the bug?

if num <= *max_heap.peek().unwrap() {
    max_heap.push(num);
} else {
    min_heap.push(num);
}

if max_heap.len() > min_heap.len() + 1 {
    let val = max_heap.pop().unwrap();
    min_heap.push(val);
} else if min_heap.len() > max_heap.len() {
    let val = min_heap.pop().unwrap();
    max_heap.push(val);
}

medium

A. Missing check for empty max_heap before peek() causing panic

B. Heaps are not balanced after insertion

C. Should pop from min_heap when max_heap is larger

D. No bug, code balances heaps correctly

Solution

1. Step 1: Analyze peek() usage on max_heap

   Calling peek() on an empty max_heap causes unwrap() panic because there is no element.

2. Step 2: Identify missing empty check

   Code should check if max_heap is empty before peek() to avoid runtime error.

3. Final Answer:

   Missing check for empty max_heap before peek() causing panic -> Option A

4. Quick Check:

   peek() on empty heap = panic without check [OK]

Hint: Always check heap is not empty before peek() [OK]

Common Mistakes:

• Assuming heaps always have elements before peek()
• Confusing which heap to pop from when balancing
• Ignoring unwrap() panic possibility

5. How would you modify the two heaps approach to efficiently find the median if the data stream can contain duplicate numbers and you want to remove numbers as well?

hard

A. Use only one heap and sort after every insertion and removal

B. Use two heaps plus a hash map to track counts and lazy removal of elements

C. Use a balanced binary search tree instead of heaps

D. Ignore duplicates and only keep unique numbers in heaps

Solution

1. Step 1: Understand problem with duplicates and removals

   Heaps alone cannot remove arbitrary elements efficiently, especially duplicates.

2. Step 2: Use hash map for counts and lazy removal

   Track counts of elements in a hash map and mark elements for removal without immediate heap deletion.

3. Step 3: Clean heaps lazily during median calculation

   Remove invalid elements from heap tops lazily to maintain correct median.

4. Final Answer:

   Use two heaps plus a hash map to track counts and lazy removal of elements -> Option B

5. Quick Check:

   Heaps + hash map = efficient median with removals [OK]

Hint: Combine heaps with hash map for removals and duplicates [OK]

Common Mistakes:

• Trying to remove elements directly from heaps
• Ignoring duplicates in median calculation
• Using only sorting which is inefficient