From c3ce8dd95e94d96273a0f4ce0f8595b5c4f99ef1 Mon Sep 17 00:00:00 2001 From: beam Date: Mon, 26 Oct 2026 10:29:44 +0000 Subject: [PATCH] beam: sliding window median with two-heap lazy deletion + tests --- wizard-work/beam/sliding_window_median.py | 108 ++++++++++++++++++ .../beam/test_sliding_window_median.py | 50 ++++++++ 2 files changed, 158 insertions(+) create mode 100644 wizard-work/beam/sliding_window_median.py create mode 100644 wizard-work/beam/test_sliding_window_median.py diff --git a/wizard-work/beam/sliding_window_median.py b/wizard-work/beam/sliding_window_median.py new file mode 100644 index 0000000..40e1c04 --- /dev/null +++ b/wizard-work/beam/sliding_window_median.py @@ -0,0 +1,108 @@ +"""Sliding-window median using two heaps with lazy deletion. + +Public API: + sliding_window_median(nums, k) -> list[float] + Median of each contiguous window of size k. For even k, median is the + average of the two middle values. Empty input / k<=0 -> []. + SlidingWindowMedian(k) -> incremental iterator-friendly helper with + .add(x), .median(). +""" +import heapq +from collections import defaultdict + + +class _DualHeap: + def __init__(self): + self.lo = [] # max-heap (negated) for lower half + self.hi = [] # min-heap for upper half + self.delayed = defaultdict(int) # value -> pending deletions + self.lo_size = 0 # valid count in lo + self.hi_size = 0 # valid count in hi + + def _prune(self, heap): + while heap: + v = -heap[0] if heap is self.lo else heap[0] + if self.delayed.get(v, 0): + heapq.heappop(heap) + self.delayed[v] -= 1 + if self.delayed[v] == 0: + del self.delayed[v] + else: + break + + def _balance(self): + # lo may have at most one more than hi + if self.lo_size > self.hi_size + 1: + self._prune(self.lo) + v = -heapq.heappop(self.lo) + self.lo_size -= 1 + heapq.heappush(self.hi, v) + self.hi_size += 1 + self._prune(self.lo) + elif self.lo_size < self.hi_size: + self._prune(self.hi) + v = heapq.heappop(self.hi) + self.hi_size -= 1 + heapq.heappush(self.lo, -v) + self.lo_size += 1 + self._prune(self.hi) + + def add(self, num): + if not self.lo or num <= -self.lo[0]: + heapq.heappush(self.lo, -num) + self.lo_size += 1 + else: + heapq.heappush(self.hi, num) + self.hi_size += 1 + self._balance() + + def remove(self, num): + self.delayed[num] += 1 + if num <= -self.lo[0]: + self.lo_size -= 1 + if num == -self.lo[0]: + self._prune(self.lo) + else: + self.hi_size -= 1 + self._balance() + + def median(self): + if self.lo_size == 0 and self.hi_size == 0: + raise ValueError("empty") + self._prune(self.lo) + self._prune(self.hi) + if self.lo_size > self.hi_size: + return -self.lo[0] + return (-self.lo[0] + self.hi[0]) / 2.0 + + +class SlidingWindowMedian: + def __init__(self, k): + if k <= 0: + raise ValueError("k must be positive") + self.k = k + self.dh = _DualHeap() + self.buf = [] + + def add(self, x): + self.dh.add(x) + self.buf.append(x) + if len(self.buf) > self.k: + old = self.buf.pop(0) + self.dh.remove(old) + return self.median() if len(self.buf) == self.k else None + + def median(self): + return self.dh.median() + + +def sliding_window_median(nums, k): + if k <= 0 or not nums: + return [] + swm = SlidingWindowMedian(k) + out = [] + for n in nums: + m = swm.add(n) + if m is not None: + out.append(m) + return out diff --git a/wizard-work/beam/test_sliding_window_median.py b/wizard-work/beam/test_sliding_window_median.py new file mode 100644 index 0000000..5fa7ba1 --- /dev/null +++ b/wizard-work/beam/test_sliding_window_median.py @@ -0,0 +1,50 @@ +import random, statistics +from sliding_window_median import sliding_window_median, SlidingWindowMedian + +def brute(nums, k): + out=[] + for i in range(len(nums)-k+1): + w=sorted(nums[i:i+k]) + if k%2: out.append(float(w[k//2])) + else: out.append((w[k//2-1]+w[k//2])/2.0) + return out + +def test_basic(): + assert sliding_window_median([1,3,-1,-3,5,3,6,7],3)==[1.0, -1.0, -1.0, 3.0, 5.0, 6.0] + +def test_even(): + assert sliding_window_median([1,2,3,4],2)==[1.5,2.5,3.5] + +def test_single(): + assert sliding_window_median([5],1)==[5.0] + +def test_empty(): + assert sliding_window_median([],3)==[] + assert sliding_window_median([1,2],0)==[] + +def test_random(): + for _ in range(300): + n=random.randint(1,40) + k=random.randint(1,n) + nums=[random.randint(-20,20) for _ in range(n)] + got=sliding_window_median(nums,k) + exp=brute(nums,k) + assert len(got)==len(exp) + for a,b in zip(got,exp): + assert abs(a-b)<1e-9, (nums,k,a,b) + +def test_incremental(): + swm=SlidingWindowMedian(3) + assert swm.add(1) is None + assert swm.add(3) is None + assert swm.add(-1)==1.0 + assert swm.add(-3)==-1.0 + +def test_duplicates_removal(): + # duplicates where removal ambiguity exists + assert sliding_window_median([2,2,2,2],2)==[2.0,2.0,2.0] + +if __name__=="__main__": + for name in [n for n in globals() if n.startswith("test_")]: + globals()[name]() + print("ALL PASS")