Median of Data Stream Using Two Heaps in DSA Python

Mental Model

We keep two heaps to split numbers into smaller half and larger half so we can quickly find the middle value.

Analogy: Imagine sorting people by height into two lines: one line for shorter people and one for taller people. The middle person is between these two lines.

Max-Heap (smaller half): [5, 3, 1]
Min-Heap (larger half): [6, 8, 9]
Median is between the top of these two heaps.

Dry Run Walkthrough

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

Goal: Find the median after each new number is added

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

Max-Heap: [5]
Min-Heap: []

Why: Start by adding first number to smaller half

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

Max-Heap: [3]
Min-Heap: [5]

Why: Keep smaller half smaller and balance heaps

Step 3: Add 8 to min-heap (larger half)

Max-Heap: [3]
Min-Heap: [5, 8]

Why: Add larger number to larger half

Step 4: Balance heaps by moving top of min-heap to max-heap

Max-Heap: [5, 3]
Min-Heap: [8]

Why: Heaps must be balanced in size

Step 5: Add 9 to min-heap

Max-Heap: [5, 3]
Min-Heap: [8, 9]

Why: Add larger number to larger half

Step 6: Add 1 to max-heap

Max-Heap: [5, 3, 1]
Min-Heap: [8, 9]

Why: Add smaller number to smaller half

Step 7: Add 6 to max-heap, then balance by moving top of max-heap to min-heap

Max-Heap: [5, 3, 1]
Min-Heap: [6, 8, 9]

Why: Balance heaps after adding new number

Result:

Max-Heap: [5, 3, 1]
Min-Heap: [6, 8, 9]
Median after each insertion: 5, 4.0, 5, 6.5, 5, 5.5

Annotated Code

DSA Python

import heapq

class MedianFinder:
    def __init__(self):
        self.small = []  # max-heap (store negative values)
        self.large = []  # min-heap

    def add_num(self, num: int) -> None:
        # Add to max-heap (small)
        heapq.heappush(self.small, -num)
        # Move largest from small to large to keep order
        if self.small and self.large and (-self.small[0]) > self.large[0]:
            val = -heapq.heappop(self.small)
            heapq.heappush(self.large, val)

        # Balance sizes
        if len(self.small) > len(self.large) + 1:
            val = -heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        elif len(self.large) > len(self.small):
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -val)

    def find_median(self) -> float:
        if len(self.small) > len(self.large):
            return float(-self.small[0])
        else:
            return (-self.small[0] + self.large[0]) / 2

# Driver code
finder = MedianFinder()
stream = [5, 3, 8, 9, 1, 6]
medians = []
for num in stream:
    finder.add_num(num)
    medians.append(finder.find_median())

print(", ".join(str(m) for m in medians))

heapq.heappush(self.small, -num)

Add new number to max-heap as negative to simulate max behavior

if self.small and self.large and (-self.small[0]) > self.large[0]:

Ensure max of smaller half is not bigger than min of larger half

val = -heapq.heappop(self.small) heapq.heappush(self.large, val)

Move top from smaller half to larger half to maintain order

if len(self.small) > len(self.large) + 1:

Balance heaps if smaller half too big

val = -heapq.heappop(self.small) heapq.heappush(self.large, val)

Move element from smaller half to larger half to balance sizes

elif len(self.large) > len(self.small):

Balance heaps if larger half too big

val = heapq.heappop(self.large) heapq.heappush(self.small, -val)

Move element from larger half to smaller half to balance sizes

if len(self.small) > len(self.large): return float(-self.small[0])

If smaller half bigger, median is top of smaller half

else: return (-self.small[0] + self.large[0]) / 2

If equal size, median is average of tops

Output

5, 4.0, 5, 6.5, 5, 5.5

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 all numbers each time (O(n log n)), this approach is faster for streaming data

Edge Cases

Empty stream

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

DSA Python

No explicit guard; user must add numbers before calling find_median

Stream with one number

Median is that number itself

DSA Python

if len(self.small) > len(self.large): return float(-self.small[0])

Stream with duplicate numbers

Duplicates are handled normally and included in median calculation

DSA Python

No special guard needed; heaps accept duplicates

When to Use This Pattern

When you need to find the median of a growing list of numbers efficiently, use two heaps to keep track of halves for quick median access.

Common Mistakes

Mistake: Not balancing the heaps after each insertion
Fix: Always check and balance heap sizes so their size difference is at most one

Mistake: Not maintaining order between heaps (max of smaller half <= min of larger half)
Fix: After insertion, move elements between heaps to maintain correct order

Summary

It finds the median of a stream of numbers by keeping two heaps for smaller and larger halves.

Use it when you want to get the median quickly 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 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