记录一下小根堆与大根堆的用法，以及 Go 上的实现方式

堆是一种特殊的树形数据结构，常用于实现优先队列。小根堆（Min-Heap）和大根堆（Max-Heap）分别用于快速获取最小值和最大值。

在 Go 语言中，标准库 container/heap 提供了对堆的支持，但它只提供了接口和基本实现框架，具体是小根堆还是大根堆由用户定义。

1. Go 语言堆接口简介

container/heap 包定义了一个 heap.Interface：

type Interface interface {
    sort.Interface
    Push(x interface{}) // 添加元素
    Pop() interface{}   // 弹出元素
}

其中sort.Interface定义了：

type Interface interface {
    Len() int
    Less(i, j int) bool
    Swap(i, j int)
}

根据 Less(i, j int) bool 的实现不同，堆可以是小根堆（Less 表示 “i 的元素 < j 的元素”）或者大根堆（Less 表示“i 的元素 > j 的元素”）。

2. 小根堆示例

package main

import (
    "container/heap"
    "fmt"
)

// 定义一个int切片类型，实现heap.Interface
type MinHeap []int

func (h MinHeap) Len() int            { return len(h) }
func (h MinHeap) Less(i, j int) bool  { return h[i] < h[j] } // 小根堆，父节点比子节点小
func (h MinHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }

func (h *MinHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *MinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func main() {
    h := &MinHeap{5, 3, 8, 1}
    heap.Init(h) // 初始化堆结构
    heap.Push(h, 2)
    fmt.Printf("min: %d\n", (*h)[0]) // 查看堆顶最小元素

    for h.Len() > 0 {
        fmt.Printf("%d ", heap.Pop(h))
    }
}

3. 大根堆示例

大根堆只需要把 Less 改成比较大即可：

package main

import (
    "container/heap"
    "fmt"
)

type MaxHeap []int

func (h MaxHeap) Len() int            { return len(h) }
func (h MaxHeap) Less(i, j int) bool  { return h[i] > h[j] } // 大根堆，父节点比子节点大
func (h MaxHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }

func (h *MaxHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *MaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func main() {
    h := &MaxHeap{5, 3, 8, 1}
    heap.Init(h) // 初始化堆结构
    heap.Push(h, 10)
    fmt.Printf("max: %d\n", (*h)[0]) // 查看堆顶最大元素

    for h.Len() > 0 {
        fmt.Printf("%d ", heap.Pop(h))
    }
}

4. 核心总结

  • 实现 heap.Interface 需要实现 Len(), Less(i,j int) bool, Swap(i,j int), Push(x interface{}), Pop() interface{}。
  • Less(i,j int) 控制堆的性质，小根堆 <，大根堆 >。
  • 使用 heap.Init() 初始化堆，heap.Push() 插入元素，heap.Pop() 弹出堆顶元素。
  • 不论是大根堆还是小根堆，其 Pop() interface{} 函数的逻辑都是一致的，就是剔除末位元素，因为算法会负责维护其性质
  • 如果这时候要实现一个 Peek() interface{} 函数，其功能是返回堆顶的元素，那么就只要返回其首位元素即可：h[0]（Len() > 0 的情况下）