菜单
文章 GO

Golang - 堆的实现

beautifulremi

beautifulremi

• 阅读本文需 6 分钟

接口要求

Go 的标准库 container/heap 定义了一个接口（heap.Interface），我们只需要让自己的数据类型满足这个接口，然后就可以使用 container/heap 包里提供的 Init, Push, Pop 等函数来对我们的数据进行堆操作。

核心就是实现 heap.Interface 接口。这个接口包含了 sort.Interface 的三个方法和它自己定义的两个方法。

heap.Interface 的定义如下：

type Interface interface {
    sort.Interface // 包含了 Len(), Less(i, j int) bool, Swap(i, j int)
    Push(x any)    // 在末尾添加元素
    Pop() any      // 从末尾移除元素
}

所以，我们总共需要为我们的类型实现 5个 方法：

  1. Len() int: 返回集合中的元素个数。
  2. Less(i, j int) bool: 这是定义最小堆或最大堆的关键。
    • 最小堆：如果 i 处的元素应该排在 j 处元素的前面（即 slice[i] < slice[j]），则返回 true。
    • 最大堆：如果 i 处的元素应该排在 j 处元素的前面（即 slice[i] > slice[j]），则返回 true。
  3. Swap(i, j int): 交换 i 和 j 位置的元素。
  4. Push(x any): 注意：这个方法只需要在你的数据集合（通常是切片）的末尾添加新元素 x。具体的“上浮”操作来维持堆属性，是由 heap.Push() 这个库函数完成的，它会在内部调用你写的这个 Push 方法。
  5. Pop() any: 注意：这个方法只需要从你的数据集合的末尾移除并返回元素。具体的“下沉”操作，是由 heap.Pop()这个库函数完成的，它会先把堆顶元素和末尾元素交换，然后再调用你写的这个 Pop 方法。

一个 int 切片的 heap 实现方法示例：

  1. 定义一个 IntMinHeap 类型，它本质上是一个 int 的切片。
package main

import (
	"container/heap"
	"fmt"
)

// IntMinHeap 是一个整数最小堆
type IntMinHeap []int
  1. 为类型实现 5 个接口方法
// Len 返回堆中元素的数量
func (h IntMinHeap) Len() int {
    return len(h)
}

// Less 决定了是最小堆还是最大堆
// 对于最小堆，如果 h[i] < h[j]，则返回 true
func (h IntMinHeap) Less(i, j int) bool {
    return h[i] < h[j]
}

// Swap 交换两个元素的位置
func (h IntMinHeap) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
}

// Push 将一个元素添加到堆的末尾（切片的末尾）
// 注意：这里的参数和方法接收者都必须是指针类型，因为我们要修改切片本身
func (h *IntMinHeap) Push(x any) {
    // any 类型是 interface{} 的别名
    // x.(*IntMinHeap) 是类型断言，我们确定 x 是一个 *IntMinHeap
    *h = append(*h, x.(int))
}

// Pop 从堆的末尾（切片的末尾）移除并返回元素
func (h *IntMinHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1] // 获取末尾的元素
    *h = old[0 : n-1] // 缩容切片
    return x
}
注意：Push 和 Pop 方法的接收者必须是指针类型 (*IntMinHeap)，因为它们需要修改切片（比如通过 append 改变切片的长度或容量）。其他三个方法 (Len, Less, Swap) 的接收者可以是值类型，但为了保持一致性，通常也都定义为指针类型。

包函数

1. 初始化堆 (heap.Init)

如果你有一个已经存在的切片，你可以用 heap.Init() 将它原地转换成一个堆。这个过程的时间复杂度是 $O(n)$，比一个一个 Push（总共 $O(n \log n)$）要高效。

// 创建一个堆实例
h := &IntMinHeap{9, 5, 2, 7, 1}

// 初始化堆
heap.Init(h)
fmt.Println("初始化后的堆:", *h) // 输出: [1 5 2 7 9] (不一定是完全排序的，但满足堆属性)

2. 添加元素 (heap.Push)

使用 heap.Push() 函数向堆中添加一个新元素。它会先调用你实现的 Push 方法将元素放到切片末尾，然后执行“上浮”操作来维持堆的性质。

// 添加一个新元素 0
heap.Push(h, 0)
fmt.Println("添加 0 后的堆:", *h) // 输出: [0 1 2 7 9 5]

3. 弹出堆顶元素 (heap.Pop)

使用 heap.Pop() 函数来移除并返回堆顶元素（对于最小堆就是最小值，最大堆就是最大值）。它会先将堆顶元素与末尾元素交换，然后执行“下沉”操作，最后调用你实现的 Pop 方法来移除末尾元素。

// 弹出堆顶（最小）元素
minElement := heap.Pop(h)
fmt.Printf("弹出的最小元素: %d\n", minElement.(int)) // 输出: 0
fmt.Println("弹出后的堆:", *h) // 输出: [1 5 2 7 9]

4. 查看堆顶元素（Peek）

标准库没有提供 Peek 函数，但因为堆顶元素总是位于底层切片的第一个位置 (index 0)，所以直接访问即可。

// 确保堆不为空
if h.Len() > 0 {
    fmt.Println("当前的堆顶元素:", (*h)[0]) // 输出: 1
}

5.更新元素（ heap.Fix(h Interface, i int)）

当你堆中某个元素的值（或优先级）发生了变化，但它的位置还在老地方时，堆的属性（父节点的值小于或等于其子节点的值）就可能被破坏了。heap.Fix 的作用就是在你手动修改了 h[i] 的值之后，重新调整这个元素的位置，以恢复整个堆的属性。

heap.Fix 会检查索引 i 处的元素。

  • 如果这个元素比它的父节点小（在最小堆中），它就会将该元素“上浮”（sift-up）。
  • 如果这个元素比它的某个子节点大，它就会将该元素“下沉”（sift-down）。 这个过程会一直持续，直到该元素找到它在堆中的正确位置。整个操作的时间复杂度是 $O(\log n)$。

6. 删除任意元素 （heap.Remove(h Interface, i int)）

heap.Pop 只能移除堆顶的元素。但有时我们需要从堆的中间位置移除一个指定的元素。heap.Remove 用来达到这个目的。它会返回被移除的那个元素。

直接删除中间的元素并移动其他元素来填补空缺，效率会很低 ($O(n)$)。heap.Remove 采用了一个更高效的技巧 ($O(\log n)$)：

  1. 它把你想要删除的元素（在索引 i 处）和堆的最后一个元素进行交换。
  2. 现在，原来在最后的那个元素被换到了索引 i 的位置，而要删除的元素被换到了最后。
  3. 然后它调用你实现的 Pop() 方法，轻松地将最后一个元素（也就是我们最初想删除的那个）移除。
  4. 最后，它对被交换到索引 i 的那个新元素调用 heap.Fix，因为这个元素很可能不在正确的位置，需要调整。
分享