Median of Data Stream Using Two Heaps in DSA Python - 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 find the middle score (median) at any time quickly, without sorting all scores every time a new number arrives.

🎯 Goal: Build a program that keeps track of the median of a stream of numbers using two heaps (a max heap and a min heap). You will add numbers one by one and then find the median 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 to the correct heap to keep the halves balanced.

Calculate the median from the two heaps after each insertion.

Print the median after all numbers are added.

💡 Why This Matters

🌍 Real World

Finding the median in a live data stream is useful in real-time analytics, like monitoring sensor data or financial transactions.

💼 Career

This technique is important for software engineers working on data processing, streaming platforms, and performance-critical applications.

1

Create two empty heaps

Create two empty lists called max_heap and min_heap to represent the max heap and min heap respectively.

DSA Python

# Create two empty heaps
# Your code here

Hint

Use empty lists to start your heaps. Python's heapq module works as a min heap, so for max heap you will store negative values later.

2

Add a list of numbers to process

Create a list called numbers with these exact values: [5, 15, 1, 3].

DSA Python

max_heap = []
min_heap = []
# Create the list of numbers
# Your code here

Hint

Just write the list exactly as shown.

3

Add numbers to heaps and balance them

Import heapq. Then, write a for loop using num to go through numbers. Inside the loop, add num to max_heap as a negative value using heapq.heappush(max_heap, -num). Then, if min_heap is not empty and the smallest value in min_heap is less than the largest in max_heap, move the root values between heaps to keep order. Finally, balance the sizes of the heaps so their sizes differ by at most 1 by moving elements between them using heapq.heappush and heapq.heappop.

DSA Python

max_heap = []
min_heap = []
numbers = [5, 15, 1, 3]
# Import heapq
# Add numbers to heaps and balance
# Your code here

Hint

Remember to push negative values to max_heap to simulate a max heap. Use conditions to keep heaps balanced and ordered.

4

Calculate and print the median

After the loop, check if max_heap and min_heap have the same length. If yes, calculate the median as the average of the roots of both heaps. If not, the median is the root of max_heap (remember to negate it). Print the median as a float.

DSA Python

import heapq

max_heap = []
min_heap = []
numbers = [5, 15, 1, 3]

for num in numbers:
    heapq.heappush(max_heap, -num)
    if min_heap and min_heap[0]  len(min_heap) + 1:
        val = -heapq.heappop(max_heap)
        heapq.heappush(min_heap, val)
    elif len(min_heap) > len(max_heap):
        val = heapq.heappop(min_heap)
        heapq.heappush(max_heap, -val)

# Calculate and print the median
# Your code here

Hint

If heaps are equal size, median is average of roots. Otherwise, median is root of max_heap (negated).

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