Sliding Window Median: Two Heaps and Lazy Deletion

"Return the median of every window of size k." It reads like a small twist on the running median, and the first half of the solution is exactly that: two heaps split the window into a small half and a large half. The twist is the other half. A window also loses a value on every slide, and a heap has no way to delete a value it cannot see.

The problem it solves

Given nums and a window size k, slide the window one step at a time and report the median each time. For an odd k it is the middle value; for an even k it is the average of the two middle values.

The obvious answer sorts every window: O(k log k) per window, O(n k log k) in total. Keeping one sorted list and using binary search to insert and delete is better, but each insert or delete still shifts up to k elements, so it is O(n k). Two heaps bring each step down to logarithmic work — as long as removal can be made cheap.

The intuition

Start with the running-median idea from two heaps. A max-heap called low holds the smaller half of the window, a min-heap called high holds the larger half, and their sizes differ by at most one. The median only ever needs the two roots.

Adding a value is easy: push it on the correct side and move one root across if the sizes drift. Removing is where the plan breaks. The value leaving the window can sit anywhere inside a heap, and a heap orders only parent against child. Finding it means a linear scan, and deleting from the middle means repairing the heap around the hole.

So do not look for it. Write the value down in a dead counter and leave it where it is. A buried dead value is harmless, because nothing below a root is ever read. The only moment it matters is when it rises to a root. Then, and only then, pop it. This is lazy deletion, and it turns every removal into one dictionary update plus a few pops that each value pays for once.

There is one new piece of bookkeeping. len(low) now includes dead entries, so it is no longer the size of the small half. Balance has to use live counts that are kept by hand: decrement the side the value belonged to at the moment it is marked dead.

Watch it run

The window holds four values from [1, 8, 3, 9, 5]. The median starts at (3 + 8) / 2 = 5.5. Then 1 leaves: it is a leaf in low, so it is struck through and left in place. 5 arrives, rises to the root of low, and the median becomes (5 + 8) / 2 = 6.5 without the dead 1 ever being touched. Finally 8 leaves while it is a root, and that is the one case where a pop happens at once.

The code

Every root is kept live as an invariant: whenever a dead value might have reached a root, prune clears it. That is what makes the comparison x <= -low[0] safe.

import heapq
from collections import defaultdict

def sliding_median(nums, k):
    low, high = [], []            # low: max-heap stored negated; high: min-heap
    dead = defaultdict(int)       # value -> copies marked for removal
    size = {"low": 0, "high": 0}  # LIVE counts; len(heap) also counts the dead

    def prune(heap, sign):        # pop dead values, but only from the root
        while heap and dead[sign * heap[0]]:
            dead[sign * heap[0]] -= 1
            heapq.heappop(heap)

    def rebalance():              # live sizes equal, or low one larger
        if size["low"] > size["high"] + 1:
            heapq.heappush(high, -heapq.heappop(low))
            size["low"] -= 1; size["high"] += 1
            prune(low, -1)
        elif size["low"] < size["high"]:
            heapq.heappush(low, -heapq.heappop(high))
            size["high"] -= 1; size["low"] += 1
            prune(high, 1)

    out = []
    for i, x in enumerate(nums):
        if not low or x <= -low[0]:
            heapq.heappush(low, -x); size["low"] += 1
        else:
            heapq.heappush(high, x); size["high"] += 1
        if i >= k:                            # the value leaving the window
            gone = nums[i - k]
            dead[gone] += 1                   # mark it; never search for it
            if gone <= -low[0]:
                size["low"] -= 1
                if gone == -low[0]: prune(low, -1)
            else:
                size["high"] -= 1
                if gone == high[0]: prune(high, 1)
        rebalance()
        if i >= k - 1:
            out.append(float(-low[0]) if k % 2 else (-low[0] + high[0]) / 2)
    return out

print(sliding_median([1, 8, 3, 9, 5], 4))              # [5.5, 6.5]
print(sliding_median([1, 3, -1, -3, 5, 3, 6, 7], 3))   # [1.0, -1.0, -1.0, 3.0, 5.0, 6.0]

A single rebalance per step is enough. One add and one removal change the live difference by at most two, and one move across fixes a difference of two.

To check it, the function is compared with the definition itself — slice every window and take statistics.median — on random arrays full of duplicates and window sizes from 1 to the whole array:

import random
from statistics import median

random.seed(8)
ok = True
for _ in range(3000):
    nums = [random.randint(-6, 6) for _ in range(random.randint(1, 25))]
    k = random.randint(1, len(nums))
    brute = [float(median(nums[i:i + k])) for i in range(len(nums) - k + 1)]
    ok &= sliding_median(nums, k) == brute
print(ok)                                              # True

Values between -6 and 6 guarantee many repeats, which is exactly where the side-of-the-split test and the dead counter would disagree if either were wrong.

The complexity

Every value is pushed once and popped at most once, whether by a rebalance or by a prune. Each of those operations costs a logarithm of the heap's current length, so the total is O(n log n). The heaps can hold dead entries beyond the k live ones, which is why the safe bound uses n. Writing O(n log k) assumes the heaps stay close to k entries, so say which bound you mean. Memory is O(n) in the worst case for the same reason.

If k is small, the sorted-list version with bisect is shorter to write and perfectly fast. The heap version pays off when k is large.

Where it goes wrong

• Balancing on len(heap). The lengths include dead values, so the halves look balanced when they are not. Swap the two live-size tests in rebalance for len(low) and len(high) and the second example above prints [1.0, -1.0, -1.0, 3.0, 3.0, 3.0] — the last two medians are read from the wrong place.
• Marking dead without adjusting a side. The dead counter and the live count are two halves of one removal. Update one without the other and the split drifts by one on every slide.
• Forgetting to prune after a move. Popping the root of low in rebalance can expose a dead value. If it is not cleared, the next comparison reads a value that is no longer in the window.
• Integer division for even windows. In languages with fixed-width integers, (a + b) / 2 can overflow and integer division truncates. Compute in floating point or as a + (b - a) / 2.

How to say it in an interview

"I keep the window split across a max-heap for the lower half and a min-heap for the upper half, balanced by live size, so the median is at the roots. A heap can't delete an arbitrary value, so when a value leaves I record it in a hash map as dead and adjust the live count of its side. I only physically pop dead values when they reach a root, because the roots are all I ever read. Each value is pushed and popped once, so it's O(n log n) overall."

Then name the invariant out loud — "both roots are always live" — because every line of the code exists to keep it true.