滑动窗口中位数双堆法Python实现遇TLE错误，求优化方案

你的代码在处理大测试用例（比如k=50000、数组长度100000）时超时，核心问题出在popNum方法里的list.remove()和heapq.heapify()操作——这两个操作的时间复杂度都是O(k)，每次滑动窗口都要执行一次，导致整体时间复杂度变成O(nk)，完全扛不住大数据量。

优化思路：延迟删除（Lazy Deletion） #

Python的heapq模块实现的堆不支持高效删除任意位置的元素，所以我们改用延迟删除策略：

• 用一个哈希表deleted记录每个元素需要被删除的次数
• 当堆顶元素是待删除元素时，才真正弹出它并减少计数，直到堆顶是有效元素
• 所有堆操作（push、balance、findMedian）前都先清理堆顶的无效元素

优化后的代码 #

import heapq

class Solution(object):
    def __init__(self):
        self.small = []  # 大顶堆，存较小的一半数（用负数模拟）
        self.large = []  # 小顶堆，存较大的一半数
        self.deleted = {}  # 记录待删除元素的计数

    def _clean_heap(self, heap, is_max_heap):
        """清理堆顶的无效元素"""
        while heap:
            val = heap[0]
            original_val = -val if is_max_heap else val
            if self.deleted.get(original_val, 0) > 0:
                # 堆顶是待删除元素，弹出并减少计数
                heapq.heappop(heap)
                self.deleted[original_val] -= 1
                if self.deleted[original_val] == 0:
                    del self.deleted[original_val]
            else:
                break

    def balance(self):
        # 先清理两个堆的堆顶
        self._clean_heap(self.small, is_max_heap=True)
        self._clean_heap(self.large, is_max_heap=False)
        
        # 调整两个堆的大小差不超过1
        if len(self.small) > len(self.large) + 1:
            self._clean_heap(self.small, is_max_heap=True)
            val = -heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        elif len(self.large) > len(self.small) + 1:
            self._clean_heap(self.large, is_max_heap=False)
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -val)

    def addNum(self, num):
        # 先加入大顶堆
        heapq.heappush(self.small, -num)
        
        # 检查堆顶元素是否需要清理，再判断是否要调整堆
        self._clean_heap(self.small, is_max_heap=True)
        self._clean_heap(self.large, is_max_heap=False)
        
        if self.large and -self.small[0] > self.large[0]:
            # small的最大值大于large的最小值，交换元素
            val = -heapq.heappop(self.small)
            heapq.heappush(self.large, val)
        
        self.balance()

    def findMedian(self):
        # 先清理堆顶
        self._clean_heap(self.small, is_max_heap=True)
        self._clean_heap(self.large, is_max_heap=False)
        
        if len(self.small) > len(self.large):
            return -self.small[0]
        elif len(self.large) > len(self.small):
            return self.large[0]
        else:
            return (-self.small[0] + self.large[0]) / 2.0

    def popNum(self, num):
        # 不直接删除，而是记录待删除计数
        self.deleted[num] = self.deleted.get(num, 0) + 1
        # 清理堆顶并平衡
        self.balance()

    def medianSlidingWindow(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[float]
        """
        self.small = []
        self.large = []
        self.deleted = {}
        
        result = []
        # 初始化第一个窗口
        for i in range(k):
            self.addNum(nums[i])
        result.append(self.findMedian())
        
        # 滑动窗口
        for i in range(k, len(nums)):
            # 移除窗口左端元素
            self.popNum(nums[i - k])
            # 添加窗口右端元素
            self.addNum(nums[i])
            # 记录中位数
            result.append(self.findMedian())
        
        return result

时间复杂度分析 #

• 每个元素最多被push到堆两次、pop两次，每次堆操作是O(log k)
• 哈希表的增删查都是O(1)
• 整体时间复杂度为O(n log k)，完全能处理k=50000、数组长度100000的测试用例

内容的提问来源于stack exchange，提问作者KevinL