Sliding Window Median with Two Heaps

Welcome to the final lesson in our deep dive on finding medians. In the past two sessions, we first built a data structure using two heaps to find the median of a data stream, and then we devised an efficient "lazy deletion" strategy to handle element removals. Now, it's time to put all the pieces together.

Today's goal is to implement a complete solution for the Sliding Window Median problem. We will combine the add, remove, and findMedian functionalities into a single, cohesive algorithm that processes an array with a moving window. This pattern is not just a classic interview problem; it has direct applications in fields like financial analysis for calculating rolling statistics on time-series data, such as stock prices.

The Full Algorithm: Integrating the Components

The most effective way to structure our solution is to encapsulate the heap management logic within a dedicated class, which we've been calling MedianFinder, and then use this class within a main function that handles the sliding window iteration.

The overall flow of the main function will be:

1. Initialize an empty list to store the results.
2. Create an instance of our MedianFinder class.
3. Bootstrap the window: Add the first k elements from the input array nums to the MedianFinder.
4. Calculate the median for this first window and add it to our results list.
5. Slide the window: Iterate from the k-th element to the end of nums. In each step:
   • Add the new element entering the window (nums[i]).
   • Remove the old element leaving the window (nums[i - k]).
   • Calculate the median of the current window and add it to the results.
6. Return the results list.

This structure cleanly separates the complex data structure logic from the straightforward iteration logic.

Implementing the Sliding Window Median

The Algo.monster article we've been using provides a complete and well-commented Python implementation that follows this exact structure. It brings together the add_num, remove_num, rebalance, and prune methods we've studied into a working solution.

By studying this code, you can see how the individual components—adding, removing, pruning, and rebalancing—work in concert. For instance, notice that rebalance is called after both add_num and remove_num to ensure the heaps remain balanced. Furthermore, prune is called within rebalance and remove_num to clear out any stale elements that might have floated to the top.

Visualizing the Output

The algorithm you just implemented is precisely what's used to generate smoothed data plots in many analytical contexts. For example, in finance, a rolling median can help reveal the underlying trend of a volatile stock price by filtering out short-term spikes.

In the chart above, the blue line could represent daily price changes of a stock. The orange line, which is the rolling median, provides a much smoother representation of the trend, making it easier for an analyst or an automated trading system to interpret. It's less affected by single-day price shocks compared to a simple rolling average.

Mastering the Implementation: Common Pitfalls

As you know from your LeetCode experience, a correct algorithm is only half the battle; a robust implementation that handles edge cases is what sets apart a good solution from a great one. The MedianFinder has a few subtle traps that are important to navigate.

Let's briefly revisit the most critical pitfall: Incorrect Heap Selection During Removal.

When removing a number, num, we need to decide if it came from the small heap or the large heap to update the correct size counter. A naive check like if num <= -self.small[0] seems correct. However, what if -self.small[0] is itself a stale element marked for deletion? You might compare num against an invalid value, leading you to decrement the wrong heap's size counter.

The robust solution is to prune the heaps before making the comparison. By calling self.prune(self.small) just before the if statement in remove_num, you guarantee that -self.small[0] is a valid, active element, making your comparison reliable. This is a fine point that demonstrates a deep understanding of the algorithm's mechanics.

Complexity Analysis

Let's quickly formalize the efficiency of our solution.

• Time Complexity: $O(n\log k)$
  For each of the $n$ elements in the input array, we perform one add_num and (for most elements) one remove_num operation.

  • add_num involves a heap push and a rebalance, taking $O(\log k)$ time.
  • remove_num involves marking an element ($O(1)$) and potentially pruning/rebalancing, which in the amortized case is $O(\log k)$.
    Therefore, the total time complexity is dominated by $n$ iterations of an $O(\log k)$ operation.

• Space Complexity: $O(k)$

  • The two heaps, small and large, will store at most k valid elements combined. Even with lazy deletion, the number of stale elements is typically bounded in interview-style problems, so the physical heap size is often considered proportional to k.
  • The delayed hash map stores elements to be removed. In the worst case, this could grow, but it's often related to k.
  • The result array will store $n-k+1$ elements. If we consider the output array, the complexity is $O(n)$. If we only consider auxiliary space, it's $O(k)$. In an interview, it's best to clarify this. Stating $O(k)$ for auxiliary space is standard.

Conclusion

Congratulations on assembling the complete sliding window median solver! You have successfully integrated the two-heaps data structure with a lazy deletion mechanism to create an efficient and robust algorithm.

Key Takeaways:

• A complete solution can be cleanly structured with a MedianFinder helper class and a main function that manages the sliding window loop.
• The core logic involves bootstrapping the first window and then iteratively sliding by adding the new element and removing the old one.
• The add_num, remove_num, rebalance, and prune methods are interconnected and must be called at the right moments to maintain the heaps' integrity.
• Handling edge cases, such as pruning before comparison during removal, is critical for a robust implementation.
• The final algorithm achieves an optimal time complexity of $O(n\log k)$.

You have now mastered a complex and versatile application of heaps. In our next lesson, we will explore another problem, IPO (Initial Public Offering), which uses the two-heaps pattern in a different context to solve a resource allocation challenge. This will further solidify your ability to recognize where this powerful pattern can be applied.