109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
"""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
|