Golang container/heap 详解

Golang container/heap 详解

container/heap 是 Go 标准库中提供的堆数据结构实现，它是一个基于接口的优先队列实现。

基本概念

堆(Heap)是一种特殊的完全二叉树，满足以下性质：
• 每个节点的值都大于等于（或小于等于）其子节点的值

• 大顶堆：父节点值 >= 子节点值

• 小顶堆：父节点值 <= 子节点值

接口定义

要使用 container/heap，需要实现以下接口：

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

核心函数

• heap.Init(h Interface) - 初始化堆

• heap.Push(h Interface, x interface{}) - 添加元素

• heap.Pop(h Interface) interface{} - 移除并返回堆顶元素

• heap.Remove(h Interface, i int) interface{} - 移除指定索引元素

• heap.Fix(h Interface, i int) - 修改元素后重新调整堆

应用场景

1. 优先队列：任务调度，高优先级任务先执行
2. Top K 问题：从大量数据中找出最大/最小的 K 个元素
3. Dijkstra 算法：图的最短路径计算
4. Huffman 编码：数据压缩算法
5. 定时器管理：按时间顺序执行定时任务
6. 负载均衡：选择负载最小的服务器

代码示例

1. 基本使用（整数最小堆）

package main

import (
	"container/heap"
	"fmt"
)

// 定义一个整数堆类型
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] }

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

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

func main() {
	h := &IntHeap{2, 1, 5}
	heap.Init(h)
	heap.Push(h, 3)
	fmt.Printf("最小值: %d\n", (*h)[0])
	for h.Len() > 0 {
		fmt.Printf("%d ", heap.Pop(h))
	}
	// 输出: 最小值: 1
	// 1 2 3 5
}

2. 优先队列示例

package main

import (
	"container/heap"
	"fmt"
)

// 任务结构体
type Task struct {
	priority int
	name     string
}

// 任务堆
type TaskHeap []Task

func (h TaskHeap) Len() int           { return len(h) }
func (h TaskHeap) Less(i, j int) bool { return h[i].priority > h[j].priority } // 大顶堆
func (h TaskHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

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

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

func main() {
	tasks := &TaskHeap{
		{priority: 3, name: "普通任务"},
		{priority: 5, name: "紧急任务"},
		{priority: 1, name: "低优先级任务"},
	}
	
	heap.Init(tasks)
	
	// 添加新任务
	heap.Push(tasks, Task{priority: 4, name: "高优先级任务"})
	
	// 按优先级顺序执行任务
	for tasks.Len() > 0 {
		task := heap.Pop(tasks).(Task)
		fmt.Printf("执行任务: %s (优先级: %d)\n", task.name, task.priority)
	}
}

3. Top K 问题

package main

import (
	"container/heap"
	"fmt"
)

// 最小堆实现，用于找最大的K个数
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 topK(nums []int, k int) []int {
	if k <= 0 || len(nums) < k {
		return nil
	}
	
	h := &MinHeap{}
	heap.Init(h)
	
	for _, num := range nums {
		if h.Len() < k {
			heap.Push(h, num)
		} else if num > (*h)[0] {
			heap.Pop(h)
			heap.Push(h, num)
		}
	}
	
	result := make([]int, k)
	for i := 0; i < k; i++ {
		result[k-1-i] = heap.Pop(h).(int)
	}
	return result
}

func main() {
	nums := []int{3, 2, 1, 5, 6, 4}
	k := 3
	fmt.Printf("最大的%d个数: %v\n", k, topK(nums, k))
	// 输出: 最大的3个数: [6 5 4]
}

性能特点

• 插入和删除操作的时间复杂度为 O(log n)

• 获取堆顶元素的时间复杂度为 O(1)

• 初始化堆的时间复杂度为 O(n)

注意事项

1. 使用前必须调用 heap.Init() 初始化堆
2. 直接修改堆中的元素后，需要调用 heap.Fix() 重新调整堆
3. 不要直接操作底层切片，始终使用 heap 包提供的方法
4. 如果需要稳定的排序（相同优先级保持顺序），可以在 Less() 方法中添加额外条件

container/heap 提供了灵活高效的堆实现，适用于各种需要优先队列或部分排序的场景。