Median of Data Stream Using Two Heaps in DSA Python - Deep Dive

Overview - Median of Data Stream Using Two Heaps

What is it?

Median of Data Stream Using Two Heaps is a method to find the middle value of a list of numbers that keeps changing as new numbers come in. Instead of sorting the entire list every time, it uses two special lists called heaps to keep track of the smaller and larger halves of the numbers. This way, it can quickly find the median after each new number is added. The median is the number in the middle when all numbers are sorted, or the average of the two middle numbers if there is an even count.

Why it matters

Without this method, finding the median in a changing list would be slow because sorting the whole list every time takes a lot of time. This method makes it fast and efficient, which is important in real-life situations like monitoring stock prices, sensor data, or live user ratings where data keeps coming in continuously. It helps systems respond quickly with accurate middle values without delays.

Where it fits

Before learning this, you should understand basic data structures like arrays and heaps (priority queues). After this, you can explore more advanced streaming algorithms, order statistics, or balanced trees that also help with dynamic data queries.

Mental Model

Core Idea

Keep two heaps to split the data into smaller and larger halves so the median is always at the top of one or both heaps.

Think of it like...

Imagine you have two boxes: one holds all the smaller numbers and the other holds all the bigger numbers. The median is like the number sitting right at the edge between these two boxes. By keeping these boxes balanced, you can quickly find the middle number without sorting everything again.

┌───────────────┐       ┌───────────────┐
│ Max-Heap (low)│       │ Min-Heap (high)│
│  (smaller half)│       │  (larger half) │
│  [7, 5, 3]    │       │  [8, 10, 12]   │
└───────┬───────┘       └───────┬───────┘
        │                       │
        └─────────Median────────┘
          (7 or average of 7 & 8)

Build-Up - 7 Steps

1

Foundation

Understanding Median Concept

Concept: Learn what median means in a list of numbers and how it changes with new data.

The median is the middle number in a sorted list. If the list has an odd number of elements, the median is the center one. If even, it is the average of the two center numbers. For example, in [1, 3, 5], median is 3; in [1, 3, 5, 7], median is (3+5)/2 = 4.

Result

You can identify the median by sorting and picking the middle value(s).

Understanding median as the middle value sets the foundation for why we need efficient ways to find it in changing data.

2

Foundation

Introduction to Heaps

Concept: Learn what heaps are and how they keep track of minimum or maximum values efficiently.

A heap is a special tree-based structure where the parent node is always larger (max-heap) or smaller (min-heap) than its children. This property lets us quickly find the largest or smallest element. For example, a max-heap always has the biggest number at the top.

Result

You can get the max or min element in O(1) time and insert or remove elements in O(log n) time.

Knowing heaps helps us understand how to quickly access the middle values without sorting the entire list.

3

Intermediate

Using Two Heaps to Split Data

🤔Before reading on: do you think one heap is enough to find median quickly, or do we need two? Commit to your answer.

Concept: Split the data into two heaps: max-heap for smaller half and min-heap for larger half to keep track of median efficiently.

We keep all smaller numbers in a max-heap and all larger numbers in a min-heap. The max-heap top is the biggest of the smaller half, and the min-heap top is the smallest of the larger half. This split helps us find the median by looking at the tops of these heaps.

Result

Median can be found by looking at the top elements of the two heaps without sorting all data.

Splitting data into two heaps balances the problem and allows constant-time median access.

4

Intermediate

Balancing Heaps After Insertions

🤔Before reading on: do you think heaps should always have the same size, or can one be bigger? Commit to your answer.

Concept: Keep the two heaps balanced so their sizes differ by at most one to maintain correct median calculation.

After adding a new number, we insert it into one heap based on its value. Then, if one heap is bigger than the other by more than one element, we move the top element from the bigger heap to the smaller heap. This keeps the heaps balanced.

Result

Heaps remain balanced, ensuring the median is always at the top of one or both heaps.

Balancing heaps is crucial to maintain the correct median position and avoid errors.

5

Intermediate

Calculating Median from Heaps

🤔Before reading on: if heaps are equal size, do you think median is one top element or average of both? Commit to your answer.

Concept: Median is either the top of the bigger heap or the average of both tops if heaps are equal in size.

If heaps have equal size, median is average of max-heap top and min-heap top. If one heap is bigger, median is the top of that heap. This works because heaps split data into halves.

Result

Median is found in O(1) time by checking heap tops.

Knowing how to calculate median from heaps completes the streaming median algorithm.

6

Advanced

Implementing Median Finder in Python

🤔Before reading on: do you think Python's heapq supports max-heap directly or do we need a workaround? Commit to your answer.

Concept: Use Python's min-heap with negated values to simulate max-heap and implement the two-heap median finder.

Python's heapq module only supports min-heap. To create a max-heap, we insert negative values. We maintain two heaps: one min-heap for larger half and one max-heap (with negated values) for smaller half. Insertions and balancing follow the rules described earlier.

Result

A working Python class that can add numbers and return median efficiently.

Understanding language-specific details like heapq's min-heap-only nature is key to practical implementation.

7

Expert

Handling Edge Cases and Performance

🤔Before reading on: do you think this method handles duplicate numbers and large data streams efficiently? Commit to your answer.

Concept: Explore how the two-heap method deals with duplicates, large data, and performance trade-offs in real systems.

Duplicates are naturally handled since heaps store all values. The method runs in O(log n) per insertion, suitable for large streams. However, memory grows with data size. For infinite streams, techniques like sliding windows or approximate medians may be needed. Also, balancing heaps carefully avoids performance degradation.

Result

Robust median calculation that works well in most real-world streaming scenarios.

Knowing limitations and edge cases prepares you for applying this method in production and adapting it when needed.

Under the Hood

The two heaps maintain a partition of the data stream into two halves: the max-heap stores the smaller half with the largest element at the root, and the min-heap stores the larger half with the smallest element at the root. When a new number arrives, it is pushed into one heap based on comparison with current median. Then heaps are rebalanced to ensure size difference is at most one. The median is then either the root of the larger heap or the average of both roots. Internally, heap operations use a binary tree structure stored as arrays, allowing O(log n) insertion and removal.

Why designed this way?

This design was chosen to avoid sorting the entire data stream after each insertion, which is costly. Using two heaps leverages the heap property to keep track of the middle values efficiently. Alternatives like balanced binary search trees exist but are more complex to implement and maintain. The two-heap method strikes a balance between simplicity and performance, making it popular in streaming data scenarios.

Data Stream -> [Insert Number]
          ↓
 ┌─────────────────────────────┐
 │ Compare with median         │
 └─────────────┬───────────────┘
               │
      ┌────────┴─────────┐
      │                  │
┌─────────────┐    ┌─────────────┐
│ Max-Heap    │    │ Min-Heap    │
│ (smaller)   │    │ (larger)    │
│ [3, 5, 7]   │    │ [8, 10, 12] │
└─────┬───────┘    └─────┬───────┘
      │                  │
      └─────Rebalance────┘
               ↓
       Median = Top(s) of Heaps

Myth Busters - 4 Common Misconceptions

Quick: Do you think the median is always the top of the max-heap? Commit yes or no.

Common Belief: The median is always the top element of the max-heap.

Reality: The median depends on the sizes of both heaps. If heaps are equal size, median is the average of both tops; if not, it is the top of the larger heap.

Why it matters: Assuming median is always max-heap top leads to incorrect median calculation when heaps are balanced.

Quick: Do you think Python's heapq supports max-heap natively? Commit yes or no.

Common Belief: Python's heapq module supports both min-heap and max-heap directly.

Reality: heapq only supports min-heap. To simulate max-heap, values must be negated before insertion.

Why it matters: Not knowing this causes bugs or inefficient code when implementing median finder in Python.

Quick: Do you think this method sorts the entire data stream every time a new number arrives? Commit yes or no.

Common Belief: The two-heap method sorts the entire data stream after each insertion.

Reality: The method never sorts the entire data. It only inserts into heaps and rebalances, which is much faster.

Why it matters: Believing sorting happens every time underestimates the efficiency and usefulness of the method.

Quick: Do you think duplicates cause errors in the two-heap median method? Commit yes or no.

Common Belief: Duplicates break the two-heap median method or cause incorrect medians.

Reality: Duplicates are handled naturally since heaps store all values including duplicates.

Why it matters: Misunderstanding this may cause unnecessary complexity or avoidance of this method.

Expert Zone

1

Balancing heaps after every insertion is critical; even a single imbalance can cause median errors.

2

Using lazy deletion or delayed removal techniques can optimize performance when removing elements in sliding window median problems.

3

Negating values for max-heap in Python can cause subtle bugs if not handled consistently, especially when comparing or printing values.

When NOT to use

This method is not ideal when memory is limited and the data stream is infinite without bounds, as heaps grow indefinitely. Alternatives like approximate quantile algorithms (e.g., t-digest) or sliding window medians with fixed size buffers are better in such cases.

Production Patterns

In real systems, this method is used for live analytics dashboards, real-time monitoring of sensor data, and financial tickers. It is often combined with windowing techniques to compute median over recent data only. Implementations optimize heap operations and handle concurrency for multi-threaded environments.

Connections

Order Statistics Tree

Both find medians dynamically but use different data structures; heaps vs balanced trees.

Understanding heaps helps grasp order statistics trees which provide more general queries but with more complexity.

Streaming Algorithms

Median of data stream is a classic streaming problem solved by this method.

Knowing this method builds foundation for learning other streaming algorithms like count-min sketch or reservoir sampling.

Real-time Traffic Control

Both require quick median or middle value calculations to adjust flows dynamically.

Seeing median calculation as a control signal in traffic systems shows how algorithms impact real-world dynamic systems.

Common Pitfalls

#1 Not balancing heaps after insertion.

Wrong approach:

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

Correct approach:

def add_num(num):
if not max_heap or num <= -max_heap[0]:
heapq.heappush(max_heap, -num)
else:
heapq.heappush(min_heap, num)
# Balance heaps
if len(max_heap) > len(min_heap) + 1:
heapq.heappush(min_heap, -heapq.heappop(max_heap))
elif len(min_heap) > len(max_heap) + 1:
heapq.heappush(max_heap, -heapq.heappop(min_heap))

Root cause:

Forgetting to rebalance heaps causes size difference >1, breaking median calculation logic.

#2 Using Python heapq as max-heap directly without negation.

Wrong approach:

heapq.heappush(max_heap, num) # expecting max-heap behavior

Correct approach:

heapq.heappush(max_heap, -num) # negate to simulate max-heap

Root cause:

heapq only supports min-heap; misunderstanding this leads to wrong heap behavior.

#3 Calculating median incorrectly when heaps are equal size.

Wrong approach:

if len(max_heap) == len(min_heap):
median = -max_heap[0] # only max-heap top

Correct approach:

if len(max_heap) == len(min_heap):
median = (-max_heap[0] + min_heap[0]) / 2

Root cause:

Ignoring that median is average of both tops when heaps are equal size.

Key Takeaways

Median of a data stream can be found efficiently by splitting data into two heaps: max-heap for smaller half and min-heap for larger half.

Balancing the heaps after each insertion ensures the median is always accessible at the top of one or both heaps.

Python's heapq module only supports min-heap, so negating values is necessary to simulate max-heap behavior.

This method avoids sorting the entire data stream repeatedly, making it suitable for real-time and large-scale data.

Understanding edge cases like duplicates and heap balancing is essential for correct and efficient median calculation.

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