Median of a Data Stream

Welcome back. In the previous lesson, we established the "why" behind the two-heaps pattern, identifying it as the ideal structure for problems that require partitioning a dataset and efficiently accessing the boundary elements, like finding a running median.

Now, we shift our focus to the "how." Today's lesson is all about implementation. Your goal is to implement the two-heaps pattern to find the median of a data stream. We will translate the conceptual model of a max-heap and a min-heap working in tandem into robust, interview-ready code. This involves mastering the logic for adding new numbers while maintaining the crucial heap invariants.

Let's begin by recalling the core structure:

• A max-heap, let's call it small_half, to store the smaller half of the numbers.
• A min-heap, let's call it large_half, to store the larger half of the numbers.

To make this system work, we must enforce two invariants at all times:

1. Value Invariant: Every number in small_half must be less than or equal to every number in large_half.
2. Size Invariant: The sizes of the two heaps must be either equal or differ by at most one.

With these rules in place, the median is always accessible in $O(1)$ time from the heaps' top elements. The main challenge lies in designing the addNum method to preserve these invariants efficiently.

The Algorithm: A Three-Step Process

A robust implementation of addNum can be broken down into a clean, three-step process for each incoming number: insertion, rebalancing values, and rebalancing sizes. A good way to conceptualize this is to think of three distinct functions or logical blocks, as suggested in high-quality software design.

Let's dive into the most elegant and common implementation strategy for adding a new number.

Let's break down the logic from the reading.

The addNum(num) Method

Instead of first deciding which heap a new number belongs to, a more foolproof approach is to always add it to one heap and then immediately balance the values. Let's say our convention is that large_half (the min-heap) can have at most one more element than small_half.

1. Push to small_half and then move to large_half: We first push the new num into small_half (the max-heap). Then, we immediately pop the largest element from small_half (which could be the num we just added) and push it into large_half. This two-step dance elegantly ensures the value invariant is always maintained. After this, we know for certain that small_half.max() <= large_half.min().

2. Rebalance Sizes: The previous step always adds one element to large_half. This might disrupt the size invariant. If large_half now has more than one extra element compared to small_half, we restore balance by moving the smallest element from large_half back to small_half.

This sequence of operations ensures both invariants hold true after every addition, with a time complexity of $O(\log n)$ due to the heap pushes and pops.

The findMedian() Method

This part is straightforward once the heaps are correctly maintained:

• If the heaps have the same size (len(small_half) == len(large_half)), the total number of elements is even. The median is the average of the two middle elements: (small_half.max() + large_half.min()) / 2.
• If the heaps have different sizes, the total is odd. By our convention, large_half will be the larger heap, so its top element is the median: large_half.min().

From Theory to Python Code

Now, let's address the practical details of implementing this in Python. As you may know, Python's heapq library only provides a min-heap. How do we create the max-heap we need for small_half?

The standard trick is to store the negative of the numbers. A min-heap operating on negative values effectively behaves like a max-heap for the original positive values. For example, min(-5, -2, -8) is -8, which corresponds to max(5, 2, 8).

This requires careful handling:

• When you push a number into the small_half (our max-heap), you push its negative.
• When you peek at or pop from small_half, you must negate the result to get the original number back.

The following video provides an excellent walkthrough of the Python implementation, paying special attention to this negation trick.

A Clean Implementation

Combining the robust logic from Algo.monster with the Python-specifics from the video, we arrive at a clean and efficient implementation. Let's look at the final code.

Solidifying Your Understanding: Common Pitfalls

Since your goal is a deep and practical understanding for interview settings, it's vital to be aware of common mistakes. A correct but fragile solution is not enough; you need a robust one.

Being mindful of these pitfalls is what separates a basic understanding from a mastery of the pattern.

Conclusion

In this lesson, you have moved from the conceptual understanding of the two-heaps pattern to a detailed, practical implementation for finding the median of a data stream. We focused on establishing and maintaining the core value and size invariants, which are the keys to a correct solution.

Key Takeaways:

• A robust addNum method can be implemented by always pushing to one heap, immediately moving the top to the other heap to maintain the value invariant, and then rebalancing sizes.
• In Python, a max-heap is simulated by storing negative values in a min-heap, which requires careful negation on insertion and retrieval.
• The findMedian method is a simple $O(1)$ operation that checks heap sizes to determine whether to return the top of the larger heap or the average of both tops.
• The overall time complexity is $O(\log n)$ for addNum and $O(1)$ for findMedian, with $O(n)$ space complexity to store the numbers.

You are now equipped to implement one of the most classic and important heap-based algorithms. In our next lesson, we will tackle a more complex variant: what if you need to remove elements as well? This will lead us to the "Sliding Window Median" problem and introduce techniques like lazy deletion to handle element removal from heaps.