Python | Leetcode Python题解之第480题滑动窗口中位数

题目：

题解：

class DualHeap:def __init__(self, k: int):# 大根堆，维护较小的一半元素，注意 python 没有大根堆，需要将所有元素取相反数并使用小根堆self.small = list()# 小根堆，维护较大的一半元素self.large = list()# 哈希表，记录「延迟删除」的元素，key 为元素，value 为需要删除的次数self.delayed = collections.Counter()self.k = k# small 和 large 当前包含的元素个数，需要扣除被「延迟删除」的元素self.smallSize = 0self.largeSize = 0# 不断地弹出 heap 的堆顶元素，并且更新哈希表def prune(self, heap: List[int]):while heap:num = heap[0]if heap is self.small:num = -numif num in self.delayed:self.delayed[num] -= 1if self.delayed[num] == 0:self.delayed.pop(num)heapq.heappop(heap)else:break# 调整 small 和 large 中的元素个数，使得二者的元素个数满足要求def makeBalance(self):if self.smallSize > self.largeSize + 1:# small 比 large 元素多 2 个heapq.heappush(self.large, -self.small[0])heapq.heappop(self.small)self.smallSize -= 1self.largeSize += 1# small 堆顶元素被移除，需要进行 pruneself.prune(self.small)elif self.smallSize < self.largeSize:# large 比 small 元素多 1 个heapq.heappush(self.small, -self.large[0])heapq.heappop(self.large)self.smallSize += 1self.largeSize -= 1# large 堆顶元素被移除，需要进行 pruneself.prune(self.large)def insert(self, num: int):if not self.small or num <= -self.small[0]:heapq.heappush(self.small, -num)self.smallSize += 1else:heapq.heappush(self.large, num)self.largeSize += 1self.makeBalance()def erase(self, num: int):self.delayed[num] += 1if num <= -self.small[0]:self.smallSize -= 1if num == -self.small[0]:self.prune(self.small)else:self.largeSize -= 1if num == self.large[0]:self.prune(self.large)self.makeBalance()def getMedian(self) -> float:return float(-self.small[0]) if self.k % 2 == 1 else (-self.small[0] + self.large[0]) / 2class Solution:def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:dh = DualHeap(k)for num in nums[:k]:dh.insert(num)ans = [dh.getMedian()]for i in range(k, len(nums)):dh.insert(nums[i])dh.erase(nums[i - k])ans.append(dh.getMedian())return ans