使用maxHeap与minHeap计算指定第k小值和的代码错误排查

我用Python写了一个程序，想用大顶堆（maxHeap）和小顶堆（minHeap）计算数组m的和，其中m[i]是A[0]到A[i]中第(i//3+1)小的值。但运行后结果不对，预期输出是-38，实际得到-27，m数组的取值和预期不符，处理元素14时出现错误。

以下是我的代码：

import heapq

def find_m(A):
    n = len(A)
    min_heap = []
    max_heap = []
    m = []
    for i in range(n):
        if len(max_heap) < i // 3 + 1:
            heapq.heappush(max_heap, -A[i])
        else:
            if A[i] < -max_heap[0]:
                heapq.heappush(min_heap, -heapq.heappop(max_heap))
                heapq.heappush(max_heap, -A[i])
            else:
                heapq.heappush(min_heap, A[i])
        m.append(-max_heap[0])
    return sum(m)


_list = [11, 12, -20, 14, -10, -8, -7, -6, -4, -2]
print(find_m(_list))

• 预期m数组：[11, 11, -20, 11, -10, -10, -8, -8, -8, -7]，总和为-38
• 实际m数组：[11, 11, -20, 14, -10, -10, -7, -7, -7, -2]，总和为-27

问题分析 #

核心错误在于堆的维护逻辑顺序颠倒。原代码先判断大顶堆（max_heap）的长度是否达标，再决定新元素的存放位置，这会导致当大顶堆长度不足时，直接将新元素加入，而忽略该元素可能不属于前k小的集合（k = i//3+1），最终堆顶元素不是正确的第k小值。

以处理第4个元素（i=3，元素14）为例：

• 此时k = 3//3 +1 =2，需要维护大顶堆存储前2小的元素，堆顶为第2小值。
• 原代码因为max_heap当前长度为1（小于2），直接将-14推入max_heap，导致max_heap存储的是[-14,20]，堆顶对应原数14，而正确的第2小值应该是11。

修复后的代码 #

import heapq

def find_m(A):
    n = len(A)
    min_heap = []
    max_heap = []
    m = []
    for i in range(n):
        k = i // 3 + 1
        num = A[i]
        # 第一步：将新元素放入合适的堆
        if not max_heap or num <= -max_heap[0]:
            heapq.heappush(max_heap, -num)
        else:
            heapq.heappush(min_heap, num)
        
        # 第二步：调整堆的大小，确保max_heap的长度恰好为k
        # 如果max_heap过长，弹出堆顶到min_heap
        while len(max_heap) > k:
            moved_num = -heapq.heappop(max_heap)
            heapq.heappush(min_heap, moved_num)
        # 如果max_heap过短，从min_heap弹出最小元素补充到max_heap
        while len(max_heap) < k and min_heap:
            moved_num = heapq.heappop(min_heap)
            heapq.heappush(max_heap, -moved_num)
        
        m.append(-max_heap[0])
    return sum(m)


_list = [11, 12, -20, 14, -10, -8, -7, -6, -4, -2]
print(find_m(_list))  # 输出-38，符合预期

修复逻辑说明 #

1. 先分配元素到对应堆：新元素如果小于等于max_heap的堆顶原数，就加入max_heap（前k小的集合），否则加入min_heap（剩余元素）。
2. 调整堆的大小：
   • 若max_heap长度超过k，弹出堆顶（前k小中最大的元素）到min_heap。
   • 若max_heap长度不足k，从min_heap弹出最小元素补充到max_heap，确保max_heap始终保存前k小的元素，堆顶即为第k小值。

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