Finding the Median from a Data Stream

1. The Challenge: A "Running" Median

First, let's define the problem. You're receiving a continuous stream of numbers, one at a time. Your data structure needs to support two operations:

• addNum(num): Add a new number to your collection.
• findMedian(): Return the median of all numbers added so far.

Recall that the median is the middle value in a sorted list.

• If the list has an odd number of elements, the median is the single middle element. Example: [2, 3, 8] -> median is 3.
• If the list has an even number of elements, the median is the average of the two middle elements. Example: [2, 3, 8, 10] -> median is (3 + 8) / 2 = 5.5.

A naive approach would be to store all numbers in a dynamic array. For every findMedian() call, you would sort the array and find the middle element(s). Given your background in complexity analysis, you can see the issue: sorting takes $O(N\log N)$ time, which is too slow if findMedian() is called frequently.

We need a more intelligent way to track the "middle" of the dataset without fully sorting it every time.

2. The Core Insight: Splitting the Stream with Two Heaps

To find the median, we only need access to the middle one or two elements. This suggests that if we could partition the numbers into two halves—a "smaller half" and a "larger half"—we could solve the problem efficiently.

This is where heaps shine. We can use two heaps to maintain this partition:

1. A Max-Heap to store the smaller half of the numbers. The top of this heap (max-heap.peek()) will give us the largest element in the smaller half.
2. A Min-Heap to store the larger half of the numbers. The top of this heap (min-heap.peek()) will give us the smallest element in the larger half.

By keeping these two heaps balanced, the median will always be either the top of one of the heaps or the average of both tops.

To solidify this core idea, let's dive into the intuition behind this design.

3. Maintaining the Heap Invariants

For this system to work, we must enforce two critical rules, or invariants, at all times:

1. Value Invariant: Every number in the max-heap (the "small" half) must be less than or equal to every number in the min-heap (the "large" half).
2. Size Invariant: The two heaps must be balanced. Their sizes can differ by at most 1.

How do we maintain these invariants every time we call addNum? The process is quite clever. Let's watch a detailed walkthrough that demonstrates the algorithm for adding numbers and rebalancing the heaps.

Let's summarize the logic for addNum(num):

1. Add to a Heap: A simple rule is to always add the new number to one of the heaps first. Let's say we always add to the max-heap (small half).
2. Maintain Value Invariant: The number we just added might be larger than some numbers in the min-heap (large half), violating our rule. To fix this, we immediately take the largest element from the max-heap (its top) and move it to the min-heap. This single transfer elegantly ensures the value invariant holds.
3. Maintain Size Invariant (Rebalance): After the transfer, one heap might have grown too large. We check the sizes. If one heap has more than one extra element compared to the other, we move the top element from the larger heap to the smaller one.

By convention, we often maintain the balance such that if there's an odd number of elements, the heap corresponding to the larger numbers (the min-heap) holds the extra element.

4. Implementation in Python

With the logic for adding and balancing clear, the findMedian operation becomes straightforward.

• If the heaps have the same size (even total elements), the median is (max_heap.top() + min_heap.top()) / 2.
• If one heap has one more element (odd total elements), the median is larger_heap.top().

Now, let's look at the implementation. A key detail in Python is that the heapq library only provides a min-heap. To simulate a max-heap, we use a clever trick: we insert the negative of the numbers. This way, the smallest negative number (which corresponds to the largest absolute value) will be at the top of the heap.

Here's a clean implementation based on the principles discussed:

import heapq

class MedianFinder:
    def __init__(self):
        # The 'small' half of numbers, stored as a max-heap
        # We store negative values to simulate a max-heap with heapq
        self.small_heap = []  
        
        # The 'large' half of numbers, stored as a min-heap
        self.large_heap = []

    def addNum(self, num: int) -> None:
        # Add to max-heap first
        heapq.heappush(self.small_heap, -1 * num)
        
        # Ensure value invariant: every num in small_heap <= every num in large_heap
        if self.small_heap and self.large_heap and (-1 * self.small_heap[0]) > self.large_heap[0]:
            val = -1 * heapq.heappop(self.small_heap)
            heapq.heappush(self.large_heap, val)
            
        # Ensure size invariant: heaps' sizes differ by at most 1
        if len(self.small_heap) > len(self.large_heap) + 1:
            val = -1 * heapq.heappop(self.small_heap)
            heapq.heappush(self.large_heap, val)
        
        if len(self.large_heap) > len(self.small_heap) + 1:
            val = heapq.heappop(self.large_heap)
            heapq.heappush(self.small_heap, -1 * val)

    def findMedian(self) -> float:
        # If small_heap is larger (odd total)
        if len(self.small_heap) > len(self.large_heap):
            return -1 * self.small_heap[0]
        
        # If large_heap is larger (odd total)
        if len(self.large_heap) > len(self.small_heap):
            return self.large_heap[0]
        
        # If sizes are equal (even total)
        return (-1 * self.small_heap[0] + self.large_heap[0]) / 2.0

Complexity Analysis

Let's formally analyze the performance.

• $N$ = total number of elements added so far.

• Time Complexity:

  • addNum: Each call involves one or two heap pushes and pops. Since each heap operation costs $O(\log N)$, the total time is $O(\log N)$.
  • findMedian: This only requires looking at the top elements of the heaps, which is a $O(1)$ operation.

• Space Complexity:

  • We need to store all $N$ numbers in the two heaps, so the space complexity is $O(N)$.

This is a highly efficient solution that perfectly balances the cost of adding numbers and finding the median, making it ideal for data stream scenarios.

Conclusion

Congratulations on completing the Heaps and Priority Queues module! You've tackled some of the most important heap-based interview patterns, from finding "Top K" elements to merging K sorted lists, and now, finding the median from a data stream.

Key Takeaways:

• The "Find Median from Data Stream" problem can be solved efficiently using two heaps: a max-heap for the smaller half and a min-heap for the larger half.
• Maintaining a value invariant (small half <= large half) and a size invariant (sizes differ by at most 1) is crucial.
• This design achieves $O(\log N)$ for addNum and $O(1)$ for findMedian, a massive improvement over naive sorting.
• In Python, you can simulate a max-heap with the heapq (min-heap) library by negating the values.

This lesson concludes our deep dive into heap-based data structures. In our next module, we will shift focus from data structures to powerful algorithmic techniques, starting with Recursion and Backtracking, where you'll learn to solve problems involving exploration and combinatorial generation.