如何基于Go实现固定容量PriorityQueue查找Top N元素

原有代码核心错误 #

• 混淆了container/heap包的调用规则：heap.Push()、heap.Pop()是包提供的全局方法，会自动完成堆结构的调整，直接调用自定义PriorityQueue结构体的Push/Pop方法只会修改底层切片，不会维护堆的有序性，这是最核心的问题
• 多余的heap.Init调用：heap.Init仅需在堆初始化阶段执行一次，每次循环都调用会严重劣化性能，时间复杂度从O(n log k)变为O(nk)，完全失去堆结构的优势
• 原有逻辑对堆的有序性维护完全缺失，才会出现错误结果

标准实现步骤 #

首先需要让自定义的优先级队列实现heap.Interface接口，包含Len()、Less()、Swap()、Push()、Pop()五个方法，因为我们要实现小顶堆（堆顶存储当前堆内最小元素，用于Top K筛选），Less()方法要返回索引i对应元素的值小于索引j对应元素的值：

import "container/heap"

type Item struct {
    value    int // 用于比较的元素值
    priority int // 业务优先级，可按需调整
}

type PriorityQueue []*Item

func (pq PriorityQueue) Len() int { return len(pq) }

// 小顶堆实现，Less返回i位置元素是否小于j位置元素
func (pq PriorityQueue) Less(i, j int) bool {
    return pq[i].value < pq[j].value
}

func (pq PriorityQueue) Swap(i, j int) {
    pq[i], pq[j] = pq[j], pq[i]
}

// Push方法供heap包内部调用，不要直接调用
func (pq *PriorityQueue) Push(x interface{}) {
    item := x.(*Item)
    *pq = append(*pq, item)
}

// Pop方法供heap包内部调用，不要直接调用
func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    old[n-1] = nil // 避免内存泄漏
    *pq = old[0 : n-1]
    return item
}

核心遍历逻辑调整如下：

func heapSearch(items map[int]int, k int) PriorityQueue {
    // 初始化容量为k的优先级队列
    pq := make(PriorityQueue, 0, k)
    // 仅初始化时调用一次heap.Init
    heap.Init(&pq)

    for value, priority := range items {
        // 堆未满 或 当前元素大于堆顶最小元素，才需要入堆
        if pq.Len() < k || value > pq[0].value {
            // 堆已满先弹出堆顶最小元素
            if pq.Len() == k {
                heap.Pop(&pq)
            }
            // 调用heap包的Push方法自动维护堆结构
            heap.Push(&pq, &Item{
                value:    value,
                priority: priority,
            })
        }
    }
    return pq
}

补充说明 #

如果你需要实现大顶堆，只需要修改Less方法的比较逻辑，返回pq[i].value > pq[j].value即可，其他逻辑保持不变。

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