Sliding Window Median
Intuition
Solution
class Solution:
def pop(self, heap, valid):
"""
:param heap: each item in heap is a pair (value, index) where index is
the index in list nums
:param valid: minimal index which is valid
"""
index = valid - 1
while index < valid:
value, index = heapq.heappop(heap)
return value, index
def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
"""
Two heap
Almost the same problem as 295. Find Median from data stream.
"""
medians = []
left = []
left_len = 0
right = []
right_len = 0
in_left = [False] * len(nums)
for index, num in enumerate(nums):
valid = index - k + 1
# Always keep right >= left
if left_len < right_len:
heapq.heappush(right, (num, index))
v, i = self.pop(right, valid)
heapq.heappush(
left,
(-v, i)
)
left_len += 1
in_left[index] = False
in_left[i] = True
else:
heapq.heappush(left, (-num, index))
v, i = self.pop(left, valid)
heapq.heappush(
right,
(-v, i)
)
right_len += 1
in_left[index] = True
in_left[i] = False
if left_len + right_len == k:
while left and left[0][1] < valid:
heapq.heappop(left)
while right and right[0][1] < valid:
heapq.heappop(right)
if left_len == right_len:
medians.append((-left[0][0] + right[0][0]) / 2)
else:
medians.append(right[0][0])
# Remove first in the heap
# Need to use lazy delete here because python does not support
# remove
if in_left[index - k + 1]:
left_len -= 1
else:
right_len -= 1
return medians