堆与双堆模板

一句话说明

单堆通常用来维护 Top K 或当前极值；双堆则把数据分成较小一半和较大一半，最典型应用就是中位数。

Go 里堆的写法先统一

Go 通常直接基于 container/heap 实现。
最小堆默认最自然，最大堆一般通过改 Less 或取反来做。

Go 模板：最小堆

type IntMinHeap []int
 
func (h IntMinHeap) Len() int           { return len(h) }
func (h IntMinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntMinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
 
func (h *IntMinHeap) Push(x any) {
    *h = append(*h, x.(int))
}
 
func (h *IntMinHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

Go 模板：Top K 最大值

思路是维护一个大小为 k 的最小堆：

• 堆里永远保存当前最大的 k 个数
• 堆顶就是这 k 个数里的最小值，也就是淘汰线

func TopKLargest(nums []int, k int) []int {
    if k <= 0 {
        return nil
    }
 
    h := &IntMinHeap{}
    heap.Init(h)
 
    for _, v := range nums {
        if h.Len() < k {
            heap.Push(h, v)
            continue
        }
        if v > (*h)[0] {
            heap.Pop(h)
            heap.Push(h, v)
        }
    }
 
    result := append([]int(nil), (*h)...)
    sort.Sort(sort.Reverse(sort.IntSlice(result)))
    return result
}

双堆到底在维护什么

双堆的核心不是“两个堆而已”，而是维护这两个不变量：

• small 存较小的一半
• large 存较大的一半

同时保证：

• small 里的所有值都不大于 large
• small 的元素个数等于 large，或者只多一个

这样中位数就会落在堆顶。

Go 模板：数据流中位数

type IntMaxHeap []int
 
func (h IntMaxHeap) Len() int           { return len(h) }
func (h IntMaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h IntMaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
 
func (h *IntMaxHeap) Push(x any) {
    *h = append(*h, x.(int))
}
 
func (h *IntMaxHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}
 
type MedianFinder struct {
    small *IntMaxHeap
    large *IntMinHeap
}
 
func NewMedianFinder() *MedianFinder {
    small := &IntMaxHeap{}
    large := &IntMinHeap{}
    heap.Init(small)
    heap.Init(large)
    return &MedianFinder{small: small, large: large}
}
 
func (mf *MedianFinder) AddNum(num int) {
    heap.Push(mf.small, num)
    heap.Push(mf.large, heap.Pop(mf.small))
 
    if mf.large.Len() > mf.small.Len() {
        heap.Push(mf.small, heap.Pop(mf.large))
    }
}
 
func (mf *MedianFinder) FindMedian() float64 {
    if mf.small.Len() > mf.large.Len() {
        return float64((*mf.small)[0])
    }
    return float64((*mf.small)[0]+(*mf.large)[0]) / 2.0
}

什么时候该想到双堆

• 数据流中位数
• 动态维护第 k 大 / 第 k 小附近的边界
• 一边维护左半最大值，一边维护右半最小值

易错点

堆与双堆模板最容易错的地方

• Go 的 Pop / Push 是给 heap 包回调用的，不是你自己直接随便改切片。
• 双堆平衡时，先保证顺序关系，再保证数量关系。
• 最大堆和最小堆不要把 Less 写反。

相关主题

• 堆
• 堆排序

---

返回：算法模板