Python滑动窗口中位数优化：双堆法解决TLE问题

import heapq

class Solution(object):
    def __init__(self):
        self.small = [] # 最大堆 (存储负值模拟)
        self.large = [] # 最小堆

    def balance(self):
        # ... 堆平衡逻辑 ...

    def addNum(self, num):
        # ... 添加元素逻辑 ...

    def findMedian(self):
        # ... 查找中位数逻辑 ...

    def popNum(self, num):
        # 性能瓶颈所在
        if num > (self.small[0] * -1):
            self.large.remove(num) # O(K) 查找并移除
            heapq.heapify(self.large) # O(K) 重新堆化
        else:
            self.small.remove(num * -1) # O(K) 查找并移除
            heapq.heapify(self.small) # O(K) 重新堆化
        self.balance()

    def medianSlidingWindow(self, nums, k):
        # ... 主逻辑 ...

import heapq

# 辅助函数：用于MaxHeap将值取反
def negate(item):
    """将 (value, index) 元组中的值取反，用于模拟最大堆。"""
    return -item[0], item[1]

class WindowHeap(object):
    """
    基础窗口堆类，支持惰性删除。
    conv: 一个转换函数，用于处理元素（例如，MaxHeap需要将值取反）。
    """
    def __init__(self, conv=lambda x: x):
        self.heap = []
        self.conv = conv  # 元素转换函数 (例如，MaxHeap的取反操作)
        self.lowindex = 0 # 窗口下限索引，用于识别已删除元素

    def peek(self):
        """
        查看堆顶元素，跳过所有已删除的元素。
        返回 (value, index) 元组或 None（如果堆为空）。
        """
        while self.heap:
            item = self.conv(self.heap[0]) # 获取原始值和索引
            if item[1] >= self.lowindex: # 如果索引在当前窗口内，则有效
                return item
            # 否则，该元素已过期，从堆中弹出并忽略
            heapq.heappop(self.heap)
        return None # 堆中没有有效元素

    def push(self, item):
        """向堆中添加一个 (value, index) 元组。"""
        heapq.heappush(self.heap, self.conv(item))

    def pop(self):
        """
        弹出堆顶元素，跳过所有已删除的元素。
        返回弹出的 (value, index) 元组。
        """
        item = self.peek() # 确保获取的是有效元素
        if item:
            heapq.heappop(self.heap)
        return item

class MinWindowHeap(WindowHeap):
    """最小窗口堆，直接存储 (value, index) 元组。"""
    def __init__(self):
        super(MinWindowHeap, self).__init__(lambda x: x) # 无需转换

class MaxWindowHeap(WindowHeap):
    """最大窗口堆，通过 negate 函数将值取反存储。"""
    def __init__(self):
        super(MaxWindowHeap, self).__init__(negate) # 使用 negate 函数

class Solution(object):
    def rebalance(self, add_val):
        """
        平衡两个堆的大小，并更新平衡因子。
        add_val: 1 表示 large 堆增加，-1 表示 small 堆增加。
        """
        self.balance += add_val
        if abs(self.balance) < 2: # 堆大小差值在 -1 到 1 之间，无需平衡
            return

        # 如果 small 堆过大，将 small 堆顶元素移到 large 堆
        if self.balance > 1: # 意味着 small 堆比 large 堆多一个元素
            self.small.push(self.large.pop()) # 注意：这里是 large.pop() 然后 push 到 small
            # 实际上是：如果 small 比 large 多 2 个，需要从 small 移一个到 large
            # 或者 large 比 small 多 2 个，需要从 large 移一个到 small
            # 这里代码的 self.balance 含义与常规理解可能不同
            # 假设 self.balance > 0 意味着 large 堆元素多， self.balance < 0 意味着 small 堆元素多
            # 原始代码逻辑是：
            # if self.balance > 1: # 意味 large 堆比 small 堆多 2 个或以上
            #     self.small.push(self.large.pop())
            # elif self.balance < -1: # 意味 small 堆比 large 堆多 2 个或以上
            #     self.large.push(self.small.pop())

            # 修正后的平衡逻辑应为：
            # 如果 small 堆比 large 堆多两个或以上元素
            if self.small.peek() and self.large.peek() and len(self.small.heap) > len(self.large.heap) + 1:
                self.large.push(self.small.pop())
            # 如果 large 堆比 small 堆多两个或以上元素
            elif self.large.peek() and self.small.peek() and len(self.large.heap) > len(self.small.heap) + 1:
                self.small.push(self.large.pop())
            # 重新计算平衡因子
            self.balance = len(self.large.heap) - len(self.small.heap) # 假设 balance 是 large - small

            # 简化平衡逻辑（根据原答案，balance 变量的更新是关键）
            # 原答案的 rebalance 逻辑是基于 self.balance 的变化来判断的
            # self.balance 初始为0，每次 insert/remove 改变其值
            # 如果 self.balance > 1，表示 large 堆比 small 堆“多”了一个元素，需要从 large 移到 small
            # 如果 self.balance < -1，表示 small 堆比 large 堆“多”了一个元素，需要从 small 移到 large
            # 这里的 self.balance 实际上记录的是 large 堆和 small 堆的“有效”元素数量差
            if self.balance > 1: # large 堆有效元素比 small 堆多 2 个或以上
                self.small.push(self.large.pop())
                self.balance -= 2 # large 减少1，small 增加1，差值减少2
            elif self.balance < -1: # small 堆有效元素比 large 堆多 2 个或以上
                self.large.push(self.small.pop())
                self.balance += 2 # small 减少1，large 增加1，差值增加2

    def insert(self, item):
        """向双堆结构中插入一个 (value, index) 元组。"""
        # 确定插入到哪个堆
        # 如果 large 堆为空，或当前 item 小于 large 堆顶，则插入 small 堆
        if not self.large.peek() or item[0] < self.large.peek()[0]:
            self.small.push(item)
            self.balance -= 1 # small 堆增加，平衡因子减1
        else:
            self.large.push(item)
            self.balance += 1 # large 堆增加，平衡因子加1
        self.rebalance(0) # 插入后进行平衡，add_val 设为 0，因为 self.balance 已经更新

    def remove(self, item):
        """
        从双堆结构中移除一个 (value, index) 元组（惰性删除）。
        通过更新 lowindex 标记元素为已删除，并重新平衡。
        """
        # 判断 item 属于哪个堆
        if item[0] >= self.large.peek()[0]: # 元素在 large 堆中
            self.balance -= 1 # large 堆有效元素减少，平衡因子减1
        else: # 元素在 small 堆中
            self.balance += 1 # small 堆有效元素减少，平衡因子加1

        # 更新 lowindex，标记该元素及其之前的所有元素为已删除
        # 注意：这里更新的是两个堆的 lowindex，确保它们只关注当前窗口内的元素
        self.large.lowindex = self.small.lowindex = item[1] + 1
        self.rebalance(0) # 移除后进行平衡

    def getMedian(self):
        """获取当前窗口的中位数。"""
        # 如果平衡因子为 0，表示两个堆有效元素数量相等
        if self.balance == 0:
            # 中位数是两个堆顶的平均值
            return (self.large.peek()[0] + self.small.peek()[0]) * 0.5
        # 如果 large 堆有效元素多，中位数是 large 堆顶
        elif self.balance > 0:
            return float(self.large.peek()[0])
        # 如果 small 堆有效元素多，中位数是 small 堆顶
        else:
            return float(self.small.peek()[0])

    def medianSlidingWindow(self, nums, k):
        """
        计算滑动窗口中位数的主函数。
        :type nums: List[int]
        :type k: int
        :rtype: List[float]
        """
        self.small = MaxWindowHeap() # 存储较小一半的元素
        self.large = MinWindowHeap() # 存储较大一半的元素
        self.balance = 0 # 平衡因子：large 堆有效元素数量 - small 堆有效元素数量

        # 将原始数组转换为 (value, index) 对列表
        items = [(val, i) for i, val in enumerate(nums)]

        # 初始化第一个窗口
        for item in items[:k]:
            self.insert(item)

        result = [self.getMedian()]

        # 滑动窗口并计算后续中位数
        # zip(items, items[k:]) 巧妙地生成 (旧元素, 新元素) 对
        for olditem, newitem in zip(items, items[k:]):
            self.remove(olditem) # 移除旧元素（惰性删除）
            self.insert(newitem) # 插入新元素
            result.append(self.getMedian())

        return result