Median in a Data Stream

1. The Challenge: Median of a Data Stream

Imagine you are receiving a continuous stream of numbers, one by one. At any point, someone can ask you for the median of all the numbers you've seen so far.

• 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. (e.g., in [1, 2, 5, 8, 11], the median is 5).
• If the list has an even number of elements, the median is the average of the two middle elements. (e.g., in [1, 2, 5, 8], the median is (2 + 5) / 2 = 3.5).

A naive approach would be to store all numbers in a list, and every time findMedian() is called, sort the list and find the middle. With N numbers, sorting takes O(N log N) time, which is too slow for a real-time stream. We need a more efficient way to add numbers and find the median.

2. The Two-Heap Strategy

The key insight is that we don't need the entire list to be sorted. We only ever need access to the middle one or two elements. To get these, we can divide the stream of numbers into two halves:

• A "small" half, containing all the numbers less than or equal to the median.
• A "large" half, containing all the numbers greater than or equal to the median.

How can we efficiently manage these two halves? We need a data structure that can give us the "boundary" elements quickly:

1. For the small half, we need its largest element. A max-heap is perfect for this.
2. For the large half, we need its smallest element. A min-heap is perfect for this.

So, our strategy is to use two heaps:

• A max-heap (small_half) to store the smaller half of the numbers.
• A min-heap (large_half) to store the larger half of the numbers.

With this setup, finding the median becomes an O(1) operation:

• If the total number of elements is odd, one heap will have one more element than the other. The median is simply the top of the larger heap.
• If the total number of elements is even, the heaps will be the same size. The median is the average of the top of the max-heap and the top of the min-heap.

The real challenge is in the addNum method: how do we add a new number while keeping the heaps balanced and ensuring that every number in small_half is less than or equal to every number in large_half?

To see a detailed walkthrough of this logic, let's watch a video that breaks down the problem and the two-heap solution.

3. Implementation in Python

As you saw in the video, the logic for addNum involves adding a new number and then potentially moving an element from one heap to the other to maintain two key invariants:

1. The Size Invariant: The heaps' sizes must differ by at most 1.
2. The Value Invariant: Every element in the small_half (max-heap) must be less than or equal to every element in the large_half (min-heap).

Now, let's look at how this is implemented in Python. A key detail is that Python's heapq module only provides a min-heap. To simulate a max-heap for our small_half, we use a simple but powerful trick: we store the negative of the numbers. This way, the smallest value in this heap (the root) corresponds to the largest actual number.

Here is a clean implementation you can study, based on the logic from the video.

Let's break down the logic from the code you just reviewed.

import heapq

class MedianFinder:
    def __init__(self):
        # self.small is a max-heap, storing the smaller half of numbers.
        # We store negative values to simulate a max-heap with Python's min-heap.
        self.small = []  
        
        # self.large is a min-heap, storing the larger half of numbers.
        self.large = []

    def addNum(self, num: int) -> None:
        # 1. Add to max-heap (small half) by default.
        heapq.heappush(self.small, -1 * num)

        # 2. Check if the value invariant is violated.
        # If the largest element in `small` is > the smallest in `large`.
        if (self.small and self.large and 
                (-1 * self.small[0]) > self.large[0]):
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)

        # 3. Rebalance the sizes if needed.
        # If `small` has more than one extra element.
        if len(self.small) > len(self.large) + 1:
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        
        # If `large` has more than one extra element.
        if len(self.large) > len(self.small) + 1:
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -1 * val)

    def findMedian(self) -> float:
        # If `small` has more elements, the median is its largest element.
        if len(self.small) > len(self.large):
            return -1 * self.small[0]
        # If `large` has more elements, the median is its smallest element.
        elif len(self.large) > len(self.small):
            return self.large[0]
        # If sizes are equal, median is the average of the two tops.
        else:
            return (-1 * self.small[0] + self.large[0]) / 2.0

4. Complexity and Pitfalls

• Time Complexity:
  • addNum: O(log N). Each call involves at most a few heap pushes and pops, each taking logarithmic time.
  • findMedian: O(1). We just look at the top element(s) of the heap(s).
• Space Complexity: O(N). We need to store all N numbers in the two heaps.

This is a highly efficient solution that meets the demands of competitive programming challenges. However, there are common mistakes to watch out for.

Conclusion

Congratulations! You've just learned a powerful and non-obvious technique for a common streaming problem. By combining two heaps, we created a data structure that can maintain the median of a dynamic dataset with very impressive time complexity.

Key Takeaways:

• Finding the median in a data stream can be solved efficiently using a max-heap for the smaller half of the numbers and a min-heap for the larger half.
• This two-heap structure allows for findMedian in O(1) time.
• The addNum operation takes O(log N) time, as it requires heap insertions/deletions to maintain the size and value invariants.
• In Python, a max-heap can be simulated using the heapq (min-heap) module by storing negated values.

Up Next:
We will continue exploring the power of heaps in our next lesson, where we'll tackle another interesting problem: "Merge K sorted arrays efficiently using a heap". This will further demonstrate how heaps can be used to manage and merge multiple sources of ordered data.