Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Go 语言标准库 —— container/heap 包（堆数据结构）

🔹 概述

container/heap 包提供了实现堆数据结构的接口和函数。

主要功能：

  • 实现最小堆或最大堆
  • 提供 Push、Pop 操作
  • 支持 Fix 操作（更新元素）
  • O(log n) 时间复杂度的插入和删除
  • O(1) 时间复杂度访问堆顶元素

重要说明：

  • heap 是一个通用的堆实现
  • 需要实现 heap.Interface 接口
  • 默认实现最小堆（可自定义为最大堆）
  • 基于切片实现
  • 不是并发安全的

堆的特性：

  • 完全二叉树结构
  • 父节点总是小于（或大于）子节点
  • 适合优先级队列、Top K 问题
  • 插入和删除的时间复杂度：O(log n)
  • 访问堆顶的时间复杂度：O(1)

应用场景：

  • 优先级队列
  • Top K 问题
  • 调度算法
  • 图算法（Dijkstra、Prim）
  • 中位数查找

🔹 接口定义

heap.Interface 接口

heap.Interface interface

  • 说明：

    • 实现堆必须实现的接口
    • 需要嵌入 sort.Interface
    • 还需要实现 Push 和 Pop 方法

  • 接口定义：

    type Interface interface {
    sort.Interface
    Push(x interface{}) // 添加元素到末尾
    Pop() interface{}   // 从末尾移除并返回元素
}

  • 需要实现的方法：

    • Len() int - 元素数量（来自 sort.Interface）
    • Less(i, j int) bool - 比较索引 i 和 j 的元素（来自 sort.Interface）
    • Swap(i, j int) - 交换索引 i 和 j 的元素（来自 sort.Interface）
    • Push(x interface{}) - 添加元素到末尾
    • Pop() interface{} - 从末尾移除并返回元素

  • 示例（最小堆）：

    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
}

🔹 核心函数

初始化堆

heap.Init(h Interface)

  • 说明：

    • 将切片初始化为堆结构
    • 时间复杂度：O(n)
    • 必须在 Push/Pop 之前调用

  • 参数：

    • h Interface - 实现了 heap.Interface 的接口

  • 注意：

    • 初始化后，h[0] 是最小（或最大）元素
    • 只需要调用一次

  • 示例：

    h := &IntHeap{5, 3, 8, 1, 9}
heap.Init(h)
fmt.Println(h[0]) // 1（最小元素）

推入元素

heap.Push(h Interface, x interface{})

  • 说明：

    • 向堆中添加一个元素
    • 自动维护堆的性质
    • 时间复杂度：O(log n)

  • 参数：

    • h Interface - 实现了 heap.Interface 的接口
    • x interface{} - 要添加的元素

  • 注意：

    • 必须先调用 heap.Init()
    • 元素会被添加到堆的正确位置

  • 示例：

    h := &IntHeap{}
heap.Init(h)
heap.Push(h, 5)
heap.Push(h, 3)
heap.Push(h, 8)
fmt.Println(h[0]) // 3（最小元素）

弹出堆顶元素

heap.Pop(h Interface) interface{}

  • 说明：

    • 移除并返回堆顶元素（最小或最大）
    • 自动维护堆的性质
    • 时间复杂度：O(log n)

  • 参数：

    • h Interface - 实现了 heap.Interface 的接口

  • 返回值：

    • interface{} - 堆顶元素

  • 注意：

    • 如果堆为空会 panic
    • 弹出后堆的大小减 1

  • 示例：

    h := &IntHeap{1, 3, 5, 8}
heap.Init(h)

min := heap.Pop(h)
fmt.Println(min) // 1

fmt.Println("剩余元素:", *h) // [3 5 8]

修复堆

heap.Fix(h Interface, i int)

  • 说明：

    • 更新索引 i 处的元素后修复堆
    • 时间复杂度：O(log n)
    • 用于元素值改变后的堆维护

  • 参数：

    • h Interface - 实现了 heap.Interface 的接口
    • i int - 需要修复的元素索引

  • 注意：

    • 元素值改变后必须调用 Fix
    • 不需要先 Pop 再 Push

  • 示例：

    h := &IntHeap{1, 3, 5, 8}
heap.Init(h)

// 修改索引 2 的元素
(*h)[2] = 0

// 修复堆
heap.Fix(h, 2)

fmt.Println(h[0]) // 0（新的最小元素）

删除指定元素

heap.Remove(h Interface, i int) interface{}

  • 说明：

    • 删除索引 i 处的元素
    • 返回被删除的元素
    • 时间复杂度：O(log n)

  • 参数：

    • h Interface - 实现了 heap.Interface 的接口
    • i int - 要删除的元素索引

  • 返回值：

    • interface{} - 被删除的元素

  • 注意：

    • 删除后堆会自动调整

  • 示例：

    h := &IntHeap{1, 3, 5, 8}
heap.Init(h)

// 删除索引 2 的元素（值为 5）
removed := heap.Remove(h, 2)

fmt.Println("删除的元素:", removed) // 5
fmt.Println("剩余元素:", *h) // [1 3 8]

🔹 实现示例

1. 最小堆（整数）

package main

import (
	"container/heap"
	"fmt"
)

// IntHeap 实现最小堆
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{}
	heap.Init(h)

	// 添加元素
	nums := []int{5, 3, 8, 1, 9, 2, 7}
	for _, n := range nums {
		heap.Push(h, n)
	}

	// 依次弹出（从小到大）
	fmt.Print("排序结果：")
	for h.Len() > 0 {
		fmt.Printf("%d ", heap.Pop(h))
	}
	// 输出：1 2 3 5 7 8 9
}

2. 最大堆（整数）

package main

import (
	"container/heap"
	"fmt"
)

// MaxHeap 实现最大堆
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{}
	heap.Init(h)

	// 添加元素
	nums := []int{5, 3, 8, 1, 9, 2, 7}
	for _, n := range nums {
		heap.Push(h, n)
	}

	// 依次弹出（从大到小）
	fmt.Print("排序结果：")
	for h.Len() > 0 {
		fmt.Printf("%d ", heap.Pop(h))
	}
	// 输出：9 8 7 5 3 2 1
}

3. 优先级队列

package main

import (
	"container/heap"
	"fmt"
)

// Item 表示优先级队列中的元素
type Item struct {
	value    string // 值
	priority int    // 优先级
	index    int    // 在堆中的索引
}

// PriorityQueue 实现优先级队列
type PriorityQueue []*Item

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]
	pq[i].index = i
	pq[j].index = j
}

func (pq *PriorityQueue) Push(x interface{}) {
	n := len(*pq)
	item := x.(*Item)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil  // 避免内存泄漏
	item.index = -1 // 标记为已移除
	*pq = old[0 : n-1]
	return item
}

// update 更新元素的优先级
func (pq *PriorityQueue) update(item *Item, value string, priority int) {
	item.value = value
	item.priority = priority
	heap.Fix(pq, item.index)
}

func main() {
	// 创建优先级队列
	pq := make(PriorityQueue, 3)
	pq[0] = &Item{value: "任务 C", priority: 3}
	pq[1] = &Item{value: "任务 A", priority: 1}
	pq[2] = &Item{value: "任务 B", priority: 2}
	
	heap.Init(&pq)

	// 添加新任务
	heap.Push(&pq, &Item{value: "任务 D", priority: 4})

	// 更新任务优先级
	pq.update(pq[1], "任务 A 升级", 5)

	// 按优先级处理任务
	fmt.Println("处理任务顺序:")
	for pq.Len() > 0 {
		item := heap.Pop(&pq).(*Item)
		fmt.Printf("%s (优先级：%d)\n", item.value, item.priority)
	}
	// 输出：
	// 任务 A 升级 (优先级：5)
	// 任务 D (优先级：4)
	// 任务 C (优先级：3)
	// 任务 B (优先级：2)
}

4. Top K 问题

package main

import (
	"container/heap"
	"fmt"
)

// IntHeap 最小堆
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
}

// topK 返回最大的 k 个数
func topK(nums []int, k int) []int {
	h := &IntHeap{}
	heap.Init(h)

	for _, num := range nums {
		heap.Push(h, num)
		// 保持堆大小为 k
		if h.Len() > k {
			heap.Pop(h)
		}
	}

	// 弹出结果（从小到大）
	result := make([]int, k)
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(h).(int)
	}
	return result
}

func main() {
	nums := []int{3, 2, 1, 5, 6, 4}
	k := 2

	result := topK(nums, k)
	fmt.Printf("Top %d: %v\n", k, result)
	// 输出：Top 2: [5 6]
}

5. 合并 K 个有序链表

package main

import (
	"container/heap"
	"fmt"
)

// ListNode 链表节点
type ListNode struct {
	Val  int
	Next *ListNode
}

// NodeHeap 最小堆
type NodeHeap []*ListNode

func (h NodeHeap) Len() int           { return len(h) }
func (h NodeHeap) Less(i, j int) bool { return h[i].Val < h[j].Val }
func (h NodeHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *NodeHeap) Push(x interface{}) {
	*h = append(*h, x.(*ListNode))
}

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

// mergeKLists 合并 K 个有序链表
func mergeKLists(lists []*ListNode) *ListNode {
	h := &NodeHeap{}
	heap.Init(h)

	// 将每个链表的头节点加入堆
	for _, node := range lists {
		if node != nil {
			heap.Push(h, node)
		}
	}

	dummy := &ListNode{}
	current := dummy

	// 依次弹出最小节点
	for h.Len() > 0 {
		node := heap.Pop(h).(*ListNode)
		current.Next = node
		current = node

		// 如果该节点有下一个节点，加入堆
		if node.Next != nil {
			heap.Push(h, node.Next)
		}
	}

	return dummy.Next
}

func main() {
	// 创建 3 个有序链表
	list1 := &ListNode{1, &ListNode{4, &ListNode{5, nil}}}
	list2 := &ListNode{1, &ListNode{3, &ListNode{4, nil}}}
	list3 := &ListNode{2, &ListNode{6, nil}}

	lists := []*ListNode{list1, list2, list3}

	// 合并
	result := mergeKLists(lists)

	// 打印结果
	for result != nil {
		fmt.Printf("%d ", result.Val)
		result = result.Next
	}
	// 输出：1 1 2 3 4 4 5 6
}

6. 数据流的中位数

package main

import (
	"container/heap"
	"fmt"
)

// MedianFinder 用于查找数据流的中位数
type MedianFinder struct {
	maxHeap *MaxHeap // 存储较小的一半
	minHeap *MinHeap // 存储较大的一半
}

// MaxHeap 最大堆
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
}

// MinHeap 最小堆
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
}

// Constructor 创建 MedianFinder
func Constructor() MedianFinder {
	return MedianFinder{
		maxHeap: &MaxHeap{},
		minHeap: &MinHeap{},
	}
}

// AddNum 添加数字
func (mf *MedianFinder) AddNum(num int) {
	// 先加入最大堆
	heap.Push(mf.maxHeap, num)
	
	// 平衡两个堆
	heap.Push(mf.minHeap, heap.Pop(mf.maxHeap))
	
	// 确保最大堆的大小 >= 最小堆
	if mf.maxHeap.Len() < mf.minHeap.Len() {
		heap.Push(mf.maxHeap, heap.Pop(mf.minHeap))
	}
}

// FindMedian 查找中位数
func (mf *MedianFinder) FindMedian() float64 {
	if mf.maxHeap.Len() == mf.minHeap.Len() {
		return float64((*mf.maxHeap)[0]+(*mf.minHeap)[0]) / 2.0
	}
	return float64((*mf.maxHeap)[0])
}

func main() {
	mf := Constructor()
	
	mf.AddNum(1)
	fmt.Printf("中位数：%.1f\n", mf.FindMedian()) // 1.0
	
	mf.AddNum(2)
	fmt.Printf("中位数：%.1f\n", mf.FindMedian()) // 1.5
	
	mf.AddNum(3)
	fmt.Printf("中位数：%.1f\n", mf.FindMedian()) // 2.0
	
	mf.AddNum(4)
	fmt.Printf("中位数：%.1f\n", mf.FindMedian()) // 2.5
}

🔹 使用场景

1. 优先级队列

// 任务调度系统
type Task struct {
	ID       int
	Priority int
	Name     string
}

// 按优先级处理任务
for pq.Len() > 0 {
	task := heap.Pop(&pq).(*Task)
	process(task)
}

2. Top K 问题

// 找出最大的 K 个数
func topK(nums []int, k int) []int {
	h := &IntHeap{}
	heap.Init(h)
	
	for _, num := range nums {
		heap.Push(h, num)
		if h.Len() > k {
			heap.Pop(h)
		}
	}
	
	result := make([]int, k)
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(h).(int)
	}
	return result
}

3. 调度算法

// CPU 调度：按优先级执行进程
for scheduler.queue.Len() > 0 {
	process := heap.Pop(&scheduler.queue).(*Process)
	execute(process)
}

4. 图算法

// Dijkstra 最短路径算法
func dijkstra(graph Graph, start int) []int {
	dist := make([]int, len(graph))
	pq := &PriorityQueue{}
	heap.Init(pq)
	
	// 初始化距离
	for i := range dist {
		dist[i] = INT_MAX
	}
	dist[start] = 0
	
	heap.Push(pq, &Item{value: start, priority: 0})
	
	for pq.Len() > 0 {
		item := heap.Pop(pq).(*Item)
		u := item.value
		
		for _, edge := range graph[u] {
			v := edge.To
			if dist[u] + edge.Weight < dist[v] {
				dist[v] = dist[u] + edge.Weight
				heap.Push(pq, &Item{value: v, priority: dist[v]})
			}
		}
	}
	
	return dist
}

🔹 注意事项和最佳实践

1. 必须实现所有接口方法

  • ✅ 必须实现 Len、Less、Swap、Push、Pop
  • ✅ Less 决定是最小堆还是最大堆
  • ⚠️ Push 和 Pop 是指针接收者

2. 初始化堆

  • ✅ 使用前必须调用 heap.Init()
  • ✅ 只需要初始化一次
  • ⚠️ 未初始化的堆行为未定义

3. 内存管理

  • ✅ Pop 后将元素置为 nil 避免内存泄漏
  • ✅ 使用指针类型时注意引用计数
func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil  // 避免内存泄漏
	*pq = old[0 : n-1]
	return item
}

4. 并发安全

  • ⚠️ heap 不是并发安全的
  • ✅ 多线程环境需要加锁
type SafeHeap struct {
	mu   sync.Mutex
	heap *IntHeap
}

func (sh *SafeHeap) Push(x int) {
	sh.mu.Lock()
	defer sh.mu.Unlock()
	heap.Push(sh.heap, x)
}

5. 性能考虑

  • ✅ 插入/删除：O(log n)
  • ✅ 访问堆顶：O(1)
  • ✅ 初始化：O(n)
  • ⚠️ 随机访问：不支持

🔥 总结

核心接口

接口/方法说明时间复杂度
heap.Interface堆接口（需实现 5 个方法）-
heap.Init(h)初始化堆O(n)
heap.Push(h, x)推入元素O(log n)
heap.Pop(h)弹出堆顶O(log n)
heap.Fix(h, i)修复堆O(log n)
heap.Remove(h, i)删除元素O(log n)

堆的类型

类型Less 实现堆顶元素
最小堆h[i] < h[j]最小值
最大堆h[i] > h[j]最大值

主要特点

  • 完全二叉树 👉 数组实现
  • 堆性质 👉 父节点 <= 子节点（最小堆）
  • 高效操作 👉 O(log n) 插入/删除
  • 快速访问 👉 O(1) 访问堆顶

使用场景

  • 优先级队列 👉 任务调度、事件处理
  • Top K 问题 👉 最大/最小 K 个数
  • 中位数查找 👉 数据流中位数
  • 图算法 👉 Dijkstra、Prim
  • 合并有序列表 👉 K 路归并

最佳实践

  • ✅ 实现完整的 heap.Interface 接口
  • ✅ 使用前调用 heap.Init()
  • ✅ 注意内存泄漏（Pop 后置 nil）
  • ✅ 根据需求选择最小堆/最大堆
  • ✅ 并发环境加锁保护
  • ⚠️ 注意：不是并发安全的

性能提示

  • 初始化 👉 O(n) 批量构建
  • 插入 👉 O(log n) heap.Push
  • 删除 👉 O(log n) heap.Pop/Remove
  • 更新 👉 O(log n) 修改后 heap.Fix
  • 访问堆顶 👉 O(1) h[0]

container/heap 包提供了高效的堆数据结构实现，适合优先级队列、Top K 问题等各种场景！