Sliding Window Median with .NET 9

#dotnet #leetcode #coding

.NET 9 introduced support for the Remove method in PriorityQueue:

boolean removed = priorityQueue.Remove(element, out var _, out var _);

This method is particularly useful for solving the Sliding Window Median problem on LeetCode. By leveraging two heaps (a max-heap and a min-heap, implemented with PriorityQueue), it allows for the immediate removal of elements that fall out of the sliding window. However, the time complexity of the Remove operation is linear, $O(n)$ .

public class Solution {
    // For a max-heap with negative values, define a custom comparer instead of using the negative priority trick
    private PriorityQueue<int, int> maxHeap = new(Comparer<int>.Create((a, b) => b.CompareTo(a)));
    private PriorityQueue<int, int> minHeap = new();

    public double[] MedianSlidingWindow(int[] nums, int k) {
        var res = new List<double>();

        for(var i = 0; i < nums.Length; i++) {
            if (i >= k) {
                Remove(nums[i-k]);
            }

            Add(nums[i]);

            if (i >= k - 1) {
                res.Add(GetMedian());
            }
        }

        return res.ToArray();
    }

    private void Add(int n) {
        // Another feature in PriorityQueue: https://learn.microsoft.com/en-us/dotnet/api/system.collections.generic.priorityqueue-2.enqueuedequeue?view=net-9.0
        var m = maxHeap.EnqueueDequeue(n, n);
        minHeap.Enqueue(m, m);

        if (minHeap.Count - 1 > maxHeap.Count)
        {
            m = minHeap.Dequeue();
            maxHeap.Enqueue(m, m);
        }
    }

    private void Remove(int n) {
        // Support for the Remove method, introduced in .NET 9
        if (!minHeap.Remove(n, out var _, out var _)) {
            maxHeap.Remove(n, out var _, out var _);
        }
    }

    private double GetMedian() {
        if (minHeap.Count != maxHeap.Count) {
            return minHeap.Peek();
        }

        // Cast to long to prevent overflow when adding values (e.g., int.MaxValue + int.MaxValue)
        return ((long)minHeap.Peek() + maxHeap.Peek()) / 2.0;
    }
}