Median of Data Stream Using Two Heaps in DSA Python - Practice Problems & Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

What is the median after inserting these numbers?

Given the sequence of numbers inserted into the data stream: 5, 15, 1, 3, what is the median after all insertions using the two heaps method?

DSA Python

from heapq import heappush, heappop

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

    def addNum(self, num):
        heappush(self.small, -num)
        if self.small and self.large and (-self.small[0]) > self.large[0]:
            val = -heappop(self.small)
            heappush(self.large, val)

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

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

finder = MedianFinder()
numbers = [5, 15, 1, 3]
for num in numbers:
    finder.addNum(num)
median = finder.findMedian()
print(median)

A 6.0

B 4.0

C 5.0

D 3.0

💡 Hint

Remember to balance the two heaps after each insertion and then calculate the median based on their sizes.

After inserting 5, 15, 1, 3, the two heaps hold values such that the median is the average of the two middle numbers: (3 + 5) / 2 = 4.0.

🧠 Conceptual

intermediate

Why use two heaps to find the median in a data stream?

What is the main reason for using two heaps (a max heap and a min heap) to find the median in a data stream?

A To keep the smaller half of numbers in max heap and larger half in min heap for quick median calculation

B To sort all numbers in ascending order efficiently

C To store all numbers twice for redundancy

D To use heaps for faster insertion only, median is calculated separately

💡 Hint

Think about how median relates to the middle values in a sorted list.

Two heaps split the data into smaller and larger halves, allowing quick access to the middle values needed for median calculation.

🔧 Debug

advanced

Identify the error in this median finder implementation

What error will this code produce when adding numbers to the MedianFinder?

DSA Python

from heapq import heappush, heappop

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

    def addNum(self, num):
        if not self.small or num <= -self.small[0]:
            heappush(self.small, -num)
        else:
            heappush(self.large, num)

        if len(self.small) > len(self.large) + 1:
            heappush(self.large, -heappop(self.small))
        elif len(self.large) > len(self.small) + 1:
            heappush(self.small, -heappop(self.large))

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

finder = MedianFinder()
numbers = [10, 20, 30]
for num in numbers:
    finder.addNum(num)
print(finder.findMedian())

A TypeError due to pushing negative value into min heap

B Correct output 20.0

C IndexError due to empty heap pop

D ValueError due to invalid comparison

💡 Hint

Check how values are pushed and popped between heaps.

The code pushes negative values into the min heap, which breaks the heap property and causes incorrect behavior or errors.

❓ Predict Output

advanced

What is the median after these insertions?

Using the two heaps method, what is the median after inserting the numbers 2, 4, 7, 1, 5, 3?

DSA Python

from heapq import heappush, heappop

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

    def addNum(self, num):
        heappush(self.small, -num)
        if self.small and self.large and (-self.small[0]) > self.large[0]:
            val = -heappop(self.small)
            heappush(self.large, val)

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

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

finder = MedianFinder()
numbers = [2, 4, 7, 1, 5, 3]
for num in numbers:
    finder.addNum(num)
median = finder.findMedian()
print(median)

A 4.0

B 3.0

C 3.5

D 5.0

💡 Hint

Balance the heaps after each insertion and then find the middle value(s).

After all insertions, the two heaps balance so the median is the average of 3 and 4, which is 3.5.

🧠 Conceptual

expert

What is the time complexity of finding median using two heaps?

What is the time complexity of adding a number and finding the median in a data stream using the two heaps approach?

A O(n) for add, O(log n) for find median

B O(log n) for add, O(log n) for find median

C O(1) for add, O(n) for find median

D O(log n) for add, O(1) for find median

💡 Hint

Consider heap insertion and peek operations.

Adding a number requires heap insertion which is O(log n), but finding median is just accessing the top elements which is O(1).

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