Sliding Window Median [Hard] — Two Heaps with Lazy Deletion

Problem Statement

Return the median array for each window of size k in the sliding window.

Input: nums=[1,3,-1,-3,5,3,6,7], k=3
Output: [1.0,-1.0,-1.0,3.0,5.0,6.0]

Intuition

Maintain two heaps:

lo: max-heap of lower half (negated for Python)
hi: min-heap of upper half

Balance: |lo| == |hi| or |lo| = |hi|+1. Median: lo[0] if odd k, else (lo[0]+hi[0])/2.

On slide: add new, mark old for lazy removal, rebalance.

Solutions

Python

import heapq
from collections import defaultdict

def medianSlidingWindow(nums: list[int], k: int) -> list[float]:
    lo = []  # max-heap (negated)
    hi = []  # min-heap
    lazy = defaultdict(int)  # lazy deletion counts
    res = []

    def push(x):
        if lo and -lo[0] >= x:
            heapq.heappush(lo, -x)
        else:
            heapq.heappush(hi, x)
        rebalance()

    def rebalance():
        while len(lo) > len(hi)+1:
            heapq.heappush(hi, -heapq.heappop(lo))
        while len(hi) > len(lo):
            heapq.heappush(lo, -heapq.heappop(hi))

    def prune(heap, sign):
        while heap and lazy[-sign*heap[0]] > 0:
            lazy[-sign*heap[0]] -= 1
            heapq.heappop(heap)

    def get_median():
        prune(lo, -1); prune(hi, 1)
        return -lo[0] if k % 2 else (-lo[0]+hi[0])/2

    for i, v in enumerate(nums):
        push(v)
        if i >= k:
            out = nums[i-k]
            lazy[out] += 1
            if out <= -lo[0]:
                prune(lo, -1)
            else:
                prune(hi, 1)
            rebalance()
        if i >= k-1:
            res.append(get_median())
    return res

Complexity

Time: O(n log k) amortized
Space: O(k)