Rolling Window Median
Problem
Slide a window of k consecutive values across a list of numbers, one position at a time from left to right, and report the median of every window. When k is even the median is the average of the two middle values. Return the medians as floats. Re-sorting each window from scratch is too slow for long lists.
Examples
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]Input: nums = [1, 2, 3, 4], k = 4
Output: [2.5]
Why: an even window averages its two middle valuesInput: nums = [5, 5, 5], k = 1
Output: [5.0, 5.0, 5.0]
Why: edge case, a window of one value is its own medianHints
0 / 3
A running median over a growing stream is a two-heap job: a max-heap for the smaller half and a min-heap for the larger half. The new difficulty is that values also have to leave.
A heap cannot remove an arbitrary value cheaply, but it does not have to remove it right away. It only matters when that value reaches the top of its heap.
Keep a count of values that have left the window but are still physically in a heap, and discard them only when they surface at a top. Track the logical size difference between the halves yourself, adjusting it on every add and every departure, and move tops across until the small half holds the same number of live values as the large half, or one more. Read the median from the live tops.
Solution
The small half lives in a max-heap and the large half in a min-heap, so the median always sits at one or both tops. When a value leaves the window it is only recorded in a pending-removal counter, and the heap it logically belongs to is decided by comparing it with the small half's top; it is physically discarded later, when it reaches a top. A balance counter tracks live sizes rather than list lengths, and tops are moved across until the small half holds as many live values as the large half, or one more. Each value is pushed and popped a constant number of times, so time is O(n log n) and space is O(n).
import heapq
from collections import Counter
def window_medians(nums, k):
lo, hi, gone, out, bal = [], [], Counter(), [], 0 # lo holds negated values
def clean(h, sign): # drop tops that already left
while h and gone[sign * h[0]]:
gone[sign * heapq.heappop(h)] -= 1
for i, x in enumerate(nums):
if lo and x <= -lo[0]: heapq.heappush(lo, -x); bal += 1
else: heapq.heappush(hi, x); bal -= 1
if i >= k: # retire nums[i - k] lazily
y = nums[i - k]; gone[y] += 1
bal += -1 if lo and y <= -lo[0] else 1
while bal > 1 or bal < 0: # live sizes must be equal or +1
clean(lo, -1); clean(hi, 1)
if bal > 1: heapq.heappush(hi, -heapq.heappop(lo)); bal -= 2
else: heapq.heappush(lo, -heapq.heappop(hi)); bal += 2
clean(lo, -1); clean(hi, 1)
if i >= k - 1:
out.append(float(-lo[0]) if k % 2 else (-lo[0] + hi[0]) / 2)
return out
print(window_medians([1, 3, -1, -3, 5, 3, 6, 7], 3)) # -> [1.0, -1.0, -1.0, 3.0, 5.0, 6.0]
print(window_medians([1, 2, 3, 4], 4)) # -> [2.5]
print(window_medians([5, 5, 5], 1)) # -> [5.0, 5.0, 5.0]Stuck on the idea rather than the code? Sliding Window Median covers it.