golang container/heap使用

container/heap 是 Go 标准库中用于实现 堆（Heap） 或 优先队列（Priority Queue） 的包。它不直接提供一个“堆类型”，而是提供一组操作函数（如 heap.Init, heap.Push, heap.Pop 等），需要自定义一个底层容器（一般是切片）并实现 heap.Interface 接口。

✅ 一、核心原理

Go 的堆是 最小堆（min-heap），即堆顶（h[0]）是最小元素。但可以通过自定义 Less 方法实现 最大堆 或其他排序逻辑。

必须实现的接口：heap.Interface

type Interface interface {
sort.Interface // 包含 Len(), Less(i, j), Swap(i, j)
Push(x interface{}) // 向堆中添加元素（注意：参数是 interface{}）
Pop() interface{} // 弹出并返回堆顶元素（返回 interface{}）
}

⚠️ 注意：Push 和 Pop 是 给 heap 包内部调用的，应使用 heap.Push(h, x) 和 heap.Pop(h) 来操作。

✅ 二、基本使用步骤

定义一个类型（一般基于 []T）。
为该类型实现 heap.Interface 的 5 个方法。
使用 heap.Init 初始化。
使用 heap.Push / heap.Pop 操作堆。

✅ 三、完整示例

示例 1：最小整数堆（Min-Heap）

package main

import (
    "container/heap"
    "fmt"
)

// 定义堆类型（基于 []int）
type IntHeap []int

// 实现 sort.Interface
func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] } // 最小堆
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

// 实现 heap.Interface 的 Push 和 Pop
func (h *IntHeap) Push(x any) {
    *h = append(*h, x.(int))
}

func (h *IntHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1]     // 取最后一个元素
    *h = old[0 : n-1] // 缩短切片
    return x
}

func main() {
    h := &IntHeap{4, 2, 5, 1, 3}

    // 初始化堆（建立堆结构）
    heap.Init(h)
    fmt.Println("初始堆:", *h) // 堆化后不必定有序，但 h[0] 是最小值

    // 入堆
    heap.Push(h, 0)
    fmt.Println("Push(0) 后:", *h)

    // 出堆（弹出最小值）
    min := heap.Pop(h).(int)
    fmt.Println("Pop() =>", min) // 0
    fmt.Println("剩余堆:", *h)
}

示例 2：最大堆（Max-Heap）

只需修改 Less 方法：

func (h IntHeap) Less(i, j int) bool {
return h[i] > h[j] // 最大堆：父节点 > 子节点
}

这样 heap.Pop() 就会弹出最大值。

示例 3：自定义结构体（如任务优先级）

package main

import (
    "container/heap"
    "fmt"
)

type Task struct {
    Name     string
    Priority int // 数字越小，优先级越高（可按需调整）
}

type TaskHeap []Task

func (h TaskHeap) Len() int      { return len(h) }
func (h TaskHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h TaskHeap) Less(i, j int) bool {
    return h[i].Priority < h[j].Priority // 最小堆：优先级数字小的在堆顶
}

func (h *TaskHeap) Push(x any) {
    *h = append(*h, x.(Task))
}

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

func main() {
    // 使用
    tasks := &TaskHeap{
       {"紧急修复", 1},
       {"日常会议", 3},
       {"代码审查", 2},
    }
    heap.Init(tasks)
    top := heap.Pop(tasks).(Task)
    fmt.Println("最高优先级任务:", top.Name) // 紧急修复
}

✅ 四、常用操作总结

操作	函数	说明
初始化堆	heap.Init(h)	将切片转换为堆结构（O(n)）
入堆	heap.Push(h, x)	插入元素并维护堆性质（O(log n)）
出堆	heap.Pop(h)	弹出堆顶元素（O(log n)）
调整堆	heap.Fix(h, i)	当 h[i] 值改变后，修复堆（O(log n)）
移除元素	heap.Remove(h, i)	移除索引 i 处的元素（O(log n)）