beam: sliding window median with two-heap lazy deletion + tests

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

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")