Median of Data Stream Using Two Heaps in DSA Python - Execution Trace

Concept Flow - Median of Data Stream Using Two Heaps

New number arrives

↓

Compare with max-heap root

↓

Add to max-heap

↓

Balance heaps size

↓

Calculate median

↓

Output median

New numbers go to one of two heaps (max-heap for smaller half, min-heap for larger half). We keep heaps balanced to find median quickly.

Execution Sample

DSA Python

max_heap = []
min_heap = []
add_num(1)
add_num(5)
add_num(3)
median = find_median()

Add numbers to two heaps and find median of the stream so far.

Execution Table

Step	Operation	Number Added	Max-Heap (smaller half)	Min-Heap (larger half)	Balance Action	Median	Visual State
1	Add number	1	[1]	[]	No balance needed	1	Max-Heap: [1], Min-Heap: []
2	Add number	5	[1]	[5]	No balance needed	3.0	Max-Heap: [1], Min-Heap: [5]
3	Add number	3	[3,1]	[5]	Balance: Move max-heap root to min-heap	3	Max-Heap: [3,1], Min-Heap: [5]
4	Find median	-	[3,1]	[5]	-	3	Median is root of max-heap since it has more elements
5	Add number	8	[3,1]	[5,8]	No balance needed	4.0	Max-Heap: [3,1], Min-Heap: [5,8]
6	Add number	9	[5,1,3]	[8,9]	Balance: Move min-heap root to max-heap	5	Max-Heap: [5,1,3], Min-Heap: [8,9]
7	Find median	-	[5,1,3]	[8,9]	-	5	Median is root of max-heap
8	Add number	2	[3,2,1]	[5,8,9]	Balance: Move max-heap root to min-heap	4.0	Max-Heap: [3,2,1], Min-Heap: [5,8,9]
9	Find median	-	[3,2,1]	[5,8,9]	-	4.0	Median is average of roots of both heaps

💡 Stream ends or no more numbers to add

Variable Tracker

Variable	Start	After 1	After 2	After 3	After 5	After 6	After 8	Final
max_heap	[]	[1]	[1]	[3,1]	[3,1]	[5,1,3]	[3,2,1]	[3,2,1]
min_heap	[]	[]	[5]	[5]	[5,8]	[8,9]	[5,8,9]	[5,8,9]
median	-	1	3.0	3	4.0	5	4.0	4.0

Key Moments - 3 Insights

Why do we add the new number to max-heap or min-heap based on comparison with max-heap root?

Because max-heap stores the smaller half of numbers, so if new number is smaller or equal to max-heap root, it belongs there; else it goes to min-heap. See execution_table steps 1-3.

Why do we balance heaps when their sizes differ by more than 1?

To keep the median calculation simple: either both heaps have equal size or max-heap has one extra element. This is shown in step 6 where balancing moves element from min-heap to max-heap.

How is median calculated from two heaps?

If heaps are equal size, median is average of roots of both heaps (step 9). If max-heap has one extra element, median is its root (steps 4 and 7).

Visual Quiz - 3 Questions

Look at the execution_table at step 3, which heap does the number 3 go into?

A Min-Heap

B Both heaps

C Max-Heap

D Neither heap

💡 Hint

Check the 'Number Added' and 'Max-Heap' columns at step 3 in execution_table

At which step does the balancing action move an element from min-heap to max-heap?

A Step 6

B Step 4

C Step 8

D Step 2

💡 Hint

Look at the 'Balance Action' column in execution_table

If we add a number smaller than all existing numbers, which heap will it be added to?

A Min-Heap

B Max-Heap

C Both heaps

D It depends on balancing

💡 Hint

Refer to key_moments about where new numbers go based on comparison with max-heap root

Concept Snapshot

Median of Data Stream Using Two Heaps:
- Use max-heap for smaller half, min-heap for larger half
- Add new number to appropriate heap based on comparison
- Balance heaps sizes to differ by at most 1
- Median is root of max-heap if odd count, else average of roots
- Efficiently find median in O(log n) per insertion

Full Transcript

This concept uses two heaps to find the median of a stream of numbers efficiently. We keep a max-heap for the smaller half of numbers and a min-heap for the larger half. When a new number arrives, we compare it with the root of the max-heap to decide which heap to add it to. After adding, we balance the heaps so their sizes differ by at most one. The median is then either the root of the max-heap if it has one extra element, or the average of the roots of both heaps if they are equal in size. This approach allows quick median calculation after each insertion.

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 keep track of the lower half and upper half of the numbers for median calculation

B. To sort all numbers in ascending order

C. To store all numbers in one data structure

D. To remove duplicates 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 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 A

4. Quick Check:

   Two heaps split data halves = efficient median [OK]

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

Common Mistakes:

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

2. Which of the following Python code snippets correctly initializes the two heaps for the Median of Data Stream problem?

easy

A. max_heap = heapq.heapify([])
min_heap = heapq.heapify([])

B. max_heap = []
min_heap = []
max_heap = heapq.heapify(max_heap)
min_heap = heapq.heapify(min_heap)

C. max_heap = [] # will store negative values for max heap behavior
min_heap = []

D. max_heap = heapq.heapify([])
min_heap = []

Solution

1. Step 1: Understand heap initialization in Python

   Python's heapq module only supports min heaps, so to simulate a max heap, we store negative values in a list.

2. Step 2: Check code correctness

   max_heap = [] # will store negative values for max heap behavior min_heap = [] correctly initializes empty lists for both heaps; max heap will store negative values to simulate max heap behavior.

3. Final Answer:

   max_heap = [] # will store negative values for max heap behavior min_heap = [] -> Option C

4. Quick Check:

   Empty lists + negative values = max heap simulation [OK]

Hint: Use negative values in list for max heap in Python [OK]

Common Mistakes:

• Calling heapq.heapify on empty list returns None
• Trying to create max heap directly
• Confusing heapq.heapify return value

3. Given the following code snippet, what will be the output after inserting the numbers 5, 15, 1, 3 in order and then calling get_median()?

import heapq

class MedianFinder:
    def __init__(self):
        self.max_heap = []
        self.min_heap = []

    def add_num(self, num):
        if not self.max_heap or num <= -self.max_heap[0]:
            heapq.heappush(self.max_heap, -num)
        else:
            heapq.heappush(self.min_heap, num)

        if len(self.max_heap) > len(self.min_heap) + 1:
            heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
        elif len(self.min_heap) > len(self.max_heap):
            heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))

    def get_median(self):
        if len(self.max_heap) > len(self.min_heap):
            return float(-self.max_heap[0])
        else:
            return (-self.max_heap[0] + self.min_heap[0]) / 2

finder = MedianFinder()
for number in [5, 15, 1, 3]:
    finder.add_num(number)
print(finder.get_median())

medium

A. 3.0

B. 6.0

C. 5.0

D. 4.0

Solution

1. Step 1: Insert numbers and balance heaps

   Insert 5: max_heap=[-5], min_heap=[]; median=5 Insert 15: max_heap=[-5], min_heap=[15]; median=(5+15)/2=10 Insert 1: max_heap=[-5,-1], min_heap=[15]; median=5 Insert 3: max_heap=[-3,-1], min_heap=[5,15]; median=(3+5)/2=4.0

2. Step 2: Calculate final median

   After all insertions, median is 4.0 as heaps are balanced and median is average of tops.

3. Final Answer:

   4.0 -> Option D

4. Quick Check:

   Median after 4 inserts = 4.0 [OK]

Hint: Balance heaps and average tops for even count [OK]

Common Mistakes:

• Not balancing heaps after insertion
• Forgetting to invert max heap values
• Calculating median incorrectly for even count

4. The following code snippet is intended to add a number to the MedianFinder class but contains a bug. What is the bug?

def add_num(self, num):
    if not self.max_heap or num < -self.max_heap[0]:
        heapq.heappush(self.max_heap, -num)
    else:
        heapq.heappush(self.min_heap, num)

    if len(self.max_heap) > len(self.min_heap):
        heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
    elif len(self.min_heap) > len(self.max_heap) + 1:
        heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))

medium

A. The balancing conditions are reversed; max_heap can grow too large

B. The comparison should be <= instead of < in the first if condition

C. The heaps are not initialized before use

D. The code uses min heap for max_heap without negation

Solution

1. Step 1: Analyze balancing conditions

   The code moves elements when max_heap size > min_heap size, but it should allow max_heap to be larger by at most 1.

2. Step 2: Identify correct balancing logic

   Correct logic: if max_heap size > min_heap size + 1, move element to min_heap; if min_heap size > max_heap size, move element to max_heap.

3. Final Answer:

   The balancing conditions are reversed; max_heap can grow too large -> Option A

4. Quick Check:

   Balance heaps with correct size difference [OK]

Hint: Max heap can be bigger by only 1 element [OK]

Common Mistakes:

• Swapping balancing conditions
• Not allowing max_heap to be bigger by one
• Incorrect comparison operators

5. You want to extend the MedianFinder to support removing numbers from the data stream while still efficiently calculating the median. Which data structure modification is best to support this feature?

hard

A. Keep using two heaps and rebuild them after each removal

B. Use balanced binary search trees (like AVL or Red-Black trees) instead of heaps

C. Use a single sorted list and binary search for insertion and removal

D. Use hash maps to store counts and ignore heaps

Solution

1. Step 1: Understand removal challenges with heaps

   Heaps do not support efficient arbitrary removal; rebuilding heaps after each removal is costly.

2. Step 2: Identify suitable data structure

   Balanced binary search trees support insertion, removal, and median retrieval in O(log n) time.

3. Step 3: Compare other options

   Sorted lists have O(n) removal; hash maps do not maintain order; rebuilding heaps is inefficient.

4. Final Answer:

   Use balanced binary search trees (like AVL or Red-Black trees) instead of heaps -> Option B

5. Quick Check:

   Balanced BSTs support fast insert/remove/median [OK]

Hint: Balanced trees allow fast insert and remove [OK]

Common Mistakes:

• Assuming heaps support fast removal
• Using sorted lists with slow removal
• Ignoring order maintenance