一上来就 nums[l:r] 排序,样例能过,提交基本就悬了。
“滑动窗口中位数”这题,坑不在中位数,坑在那个窗口一直往前滑。你每滑一步,都要删一个旧数,再塞一个新数。这个时候还想着每次把窗口整体排一遍,复杂度直接往 O(n * k log k) 走,窗口一大,时间就开始发脾气了。
这题我第一眼一般不太信“排序一次反复用”的思路。因为窗口不是只加不减,它是边进边出。能扛住这种操作的数据结构,通常就两类:平衡树,或者双堆。Python 标准库没现成平衡树,那就老老实实上双堆,再配一个延迟删除表。
思路不复杂,但细节挺脏。
我会准备两个堆:
这样一来:
麻烦在“删除旧元素”。
堆这种东西,插入快,弹堆顶也快,但你要删一个“藏在堆中间”的值,它就不配合了。所以这里不能硬删,只能记账:某个值该删了,先记到 delayed 里,等它哪天冒到堆顶,再顺手清掉。这就是延迟删除。
代码我按这个路子写,别搞得太花:
import heapqfrom collections import defaultdictclassDualHeap:def__init__(self, k: int): self.small = [] # max heap: store negative values self.large = [] # min heap self.delayed = defaultdict(int) self.small_size = 0 self.large_size = 0 self.k = kdefprune_small(self):while self.small and self.delayed[-self.small[0]]: num = -heapq.heappop(self.small) self.delayed[num] -= 1if self.delayed[num] == 0:del self.delayed[num]defprune_large(self):while self.large and self.delayed[self.large[0]]: num = heapq.heappop(self.large) self.delayed[num] -= 1if self.delayed[num] == 0:del self.delayed[num]defrebalance(self):if self.small_size > self.large_size + 1: heapq.heappush(self.large, -heapq.heappop(self.small)) self.small_size -= 1 self.large_size += 1 self.prune_small()elif self.small_size < self.large_size: heapq.heappush(self.small, -heapq.heappop(self.large)) self.small_size += 1 self.large_size -= 1 self.prune_large()defadd(self, num: int):ifnot self.small or num <= -self.small[0]: heapq.heappush(self.small, -num) self.small_size += 1else: heapq.heappush(self.large, num) self.large_size += 1 self.rebalance()defremove(self, num: int): self.delayed[num] += 1if num <= -self.small[0]: self.small_size -= 1if num == -self.small[0]: self.prune_small()else: self.large_size -= 1if self.large and num == self.large[0]: self.prune_large() self.rebalance()defmedian(self) -> float:if self.k % 2 == 1:return float(-self.small[0])return (-self.small[0] + self.large[0]) / 2.0defmedianSlidingWindow(nums, k): dh = DualHeap(k) ans = []for i in range(k): dh.add(nums[i]) ans.append(dh.median())for i in range(k, len(nums)): dh.add(nums[i]) dh.remove(nums[i - k]) ans.append(dh.median())return ans
这题真正容易写错的,不是求中位数那一行,而是平衡堆和清理脏数据的顺序。
比如旧元素刚好就在堆顶,你不先 prune,后面取中位数就会拿到一个已经出窗口的值。再比如 small_size 和 large_size 维护错了,表面上堆里有数,实际上“有效元素个数”已经乱了,最后结果会一位一位地偏。
所以这题别死盯着堆里到底存了什么,盯“有效元素”更靠谱。
说到底,这题考的不是中位数,考的是你能不能把“动态有序集合”这件事拆开处理:插入归堆,删除记账,查询只看堆顶。思路一旦顺了,这题就没那么吓人了。
这种写法的时间复杂度是 O(n log k),能过大数据。窗口题做到后面,基本都会走到这一步:不是会不会滑,而是滑的时候,手里的结构还能不能稳住。