Golang heap堆操作与实现方法

使用container/heap需实现heap.Interface接口，包括Len、Less、Swap、Push、Pop五个方法，通过heap.Init初始化堆，再用heap.Push和heap.Pop操作数据；例如IntHeap可构建小顶堆管理整数，PriorityQueue可按优先级处理任务，其中Less决定堆序，Pop从末尾移除元素，更新元素时调用heap.Fix维护结构，适用于优先队列、调度器等场景。

Go语言标准库中的container/heap提供了一个堆的接口，但它本身不直接实现堆结构，而是要求用户实现heap.Interface接口后，再通过heap.Init、heap.Push、heap.Pop等函数进行堆操作。下面详细介绍如何使用它来管理堆。

实现heap.Interface接口

要使用container/heap，必须定义一个类型并实现heap.Interface接口的五个方法：Len()、Less(i, j int)、Swap(i, j int)、Push(x)、Pop()。

定义数据类型：

type IntHeap []int

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] }

// Push 会被 heap.Push 调用，元素已由 heap.Push 添加到末尾
func (h *IntHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

// Pop 会被 heap.Pop 调用，移除并返回最小元素（根）
func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

初始化和基本操作

在使用前需要调用heap.Init将普通切片初始化为堆，之后使用heap.Push和heap.Pop进行插入和删除。

示例代码：

package main

import (
    "container/heap"
    "fmt"
)

func main() {
    h := &IntHeap{3, 1, 4, 1, 5}
    heap.Init(h)             // 初始化为堆
    heap.Push(h, 2)          // 插入元素
    fmt.Printf("最小值: %d\n", (*h)[0]) // 查看堆顶

    for h.Len() > 0 {
        min := heap.Pop(h)   // 弹出最小值
        fmt.Printf("%d ", min)
    }
    // 输出: 1 1 2 3 4 5
}

自定义结构体堆

实际开发中常需对结构体排序。比如按优先级处理任务：

type Task struct {
    ID       int
    Priority int
}

type PriorityQueue []*Task

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

func (pq PriorityQueue) Less(i, j int) bool {
    return pq[i].Priority < pq[j].Priority // 优先级小的先出（小顶堆）
}

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

func (pq *PriorityQueue) Push(x interface{}) {
    task := x.(*Task)
    *pq = append(*pq, task)
}

func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    task := old[n-1]
    *pq = old[0 : n-1]
    return task
}

使用方式：

pq := make(PriorityQueue, 0)
heap.Init(&pq)
heap.Push(&pq, &Task{ID: 1, Priority: 3})
heap.Push(&pq, &Task{ID: 2, Priority: 1})

for pq.Len() > 0 {
    task := heap.Pop(&pq).(*Task)
    fmt.Printf("执行任务 %d, 优先级 %d\n", task.ID, task.Priority)
}

注意事项与技巧

使用container/heap时注意以下几点：

• 堆数据必须是指针类型，因为Push和Pop会修改切片长度
• Less方法决定堆序性质：返回true表示i应排在j前面
• Pop总是从末尾取值，因此内部逻辑依赖Init或Fix维护结构
• 若想更新堆中元素，可修改后调用heap.Fix(&pq, index)重新调整位置

基本上就这些。掌握接口实现和标准调用流程后，就能灵活用于优先队列、调度器、TopK问题等场景。