本周我们来阅读 Go 标准库中的数据结构 heap 包，heap 相对于 sort 较为简单，相信之前在讲解 heapSort 时大家对 heap 已经有了一个初步的了解，并且 heapSort 中 siftDown 函数的逻辑与 heap 包中的基本一致，除此之外 heap 包中的 Interface 继承了 sort 中的 Interface，这几点都会使我们阅读 heap 包时更加轻松。
老样子，我们先从 Interface 的定义开始，源码如下，正如之前所说，L5 引用了 sort 包中的 Interface 这里我们不再详细说明，不了解的朋友可以去看之前的 Sort（一）文章，接下来就是 Push 和 Pop 方法，这两个方法需要使用者自己实现，但是必须保证 Push 在 Len 位置（index 从 0 开始）添加一个元素，Pop 弹出最后一个元素（Len - 1），而不能直接在堆顶操作，这样限制的具体原因与 heap 包中的 Push 与 Pop 方法实现的方式有关，稍后我们来逐一阅读。
// Note that Push and Pop in this interface are for package heap's// implementation to call. To add and remove things from the heap,// use heap.Push and heap.Pop.type Interface interface {sort.InterfacePush(x any) // add x as element Len()Pop() any // remove and return element Len() - 1.}
Heap 包中的代码不多，一共 119 行，1 个 interface 与 5 个方法和 2 个函数，对于这种结构简单的代码我们可以直接画图看一看它们之间的关系，以便更流畅的阅读源码。
Down && Up
上图（图 1）中我们可以清晰的看到，heap 包中的五个方法都是围绕 down 与 up 两个函数来进行构建的，所以我们先来阅读 down 与 up 函数。首先是 down 函数，为了更好的理解它，我们把它和 sort 包中的 siftDown 函数逐行进行对比。
先看 L2 与 L22 的两个函数签名，可以看到 siftDown 因为要对数据进行分区间的处理入参多一个 first 来标明位置，前三个参数意义基本相同，从做往右依次为要处理的对象（data、heap）、从哪开始、到那结束。L3 ～ L5 与 L23 ～ L25 完全相同，都是保存一个根节点，然后找到它的左子节点。
来到 L6 与 L26 同样都是判断子节点是否存在，但是 down 函数中多了一个条件为 j1 < 0 防溢出处理，这是因为在 L25 (2 * i + 1) 的运算过程中如果 i 值过大可能导致整数溢出变为负数，而 siftDown 中则在使用时因为 lo 的计算方式保证了寻找子节点时不会溢出，有兴趣的朋友可以回顾一下 sort 包中的 heapSort 函数。
L9 ～ L15 与 L29 ～ L36 同样都是计算右子节点的下标，然后判断下标在不在范围内，之后比较两个节点的值，然后选出一个节点与它们的父节点比较，之后按比较结果来决定是否与父节点进行交换，两个方法唯一的区别就是 siftDown 要找出比父节点大的子节点，down 要找出比父节点小的子节点，换句话说，这就是构建最大堆与最小堆方法上的唯一区别。L16 与 L37 相同，对被交换的子树进行迭代处理，最后 down 返回一个 bool 值来判断此次 down 是否对堆进行了操作，下方图 2 展示了 down 操作的过程。
// package sortfunc siftDown(data Interface, lo, hi, first int) {root := lofor {child := 2*root + 1if child >= hi {break}if child+1 < hi && data.Less(first+child, first+child+1) {child++}if !data.Less(first+root, first+child) {return}data.Swap(first+root, first+child)root = child}}// package heapfunc down(h Interface, i0, n int) bool {i := i0for {j1 := 2*i + 1if j1 >= n || j1 < 0 { // j1 < 0 after int overflowbreak}j := j1 // left childif j2 := j1 + 1; j2 < n && h.Less(j2, j1) {j = j2 // = 2*i + 2 // right child}if !h.Less(j, i) {break}h.Swap(i, j)i = j}return i > i0}
接下来是 up 函数，相对于 down 函数 up 较为简单，其主要逻辑为与父节点比较，如果满足要求，则交换，然后让父节点与父节点的父节点比较，逐级上升，直到达到到达堆顶或者不满足交换条件之后退出循环，图 3 展示了 up 操作的过程。现在就不难理解 down 与 up 作为 heap 操作的底层函数的意义了，前者从上到下构建堆，后者则是从下到上构建，下面我们从 Init 开始来分别了解 heap 的 5 个方法。
func up(h Interface, j int) {for {i := (j - 1) / 2 // parentif i == j || !h.Less(j, i) {break}h.Swap(i, j)j = i}}
Init
Init 顾名思义，初始化一个堆，在这里是最小堆，在 L8 中使用 n / 2 - 1 来寻找最后一个节点的父节点，但是这和我们之前了解的父节点计算方式 (n - 1) / 2 不同，所以注意看 L7 中 n 的定义，n 为 Len 也就是说 L8 中的 n / 2 - 1 可以理解为 (n + 1) / 2 - 1 而此式经过简单的计算是等于 (n - 1) / 2 的。之后通过循环从下到上依次对每个节点 执行 down 操作从而构建一个最小堆，图 4 展示了 Init 方法的执行过程。
// Init establishes the heap invariants required by the other routines in this package.// Init is idempotent with respect to the heap invariants// and may be called whenever the heap invariants may have been invalidated.// The complexity is O(n) where n = h.Len().func Init(h Interface) {// heapifyn := h.Len()for i := n/2 - 1; i >= 0; i-- {down(h, i, n)}}
Push && Pop
接下来是 Push 方法，L4 中使用调用者提供的 Push 方法在 Len 位置添加一个节点，之后使用 up 函数从最后一个元素开始，向上对堆进行调整，使堆重新满足最小堆，开头我们提到过，Push 必须在 Len 位置添加元素也是这个原因，如果在头部添加，则进入 up 后会直接满足退出条件不会对堆进行调整，在其它地方添加也是同理，不过还要注意另外一点，这里 L5 传入的是 Len - 1 所以我们在自己定义 Push 时也不要忘记修改堆的长度，否则新添加的节点依然不会调整到正确位置，图 5 展示了 Push 方法的执行过程。
// Push pushes the element x onto the heap.// The complexity is O(log n) where n = h.Len().func Push(h Interface, x any) {h.Push(x)up(h, h.Len()-1)}
Pop 方法，弹出堆中最小元素，即堆顶元素，首先确定最后一个元素的索引，然后交换堆顶与最后一个节点的值，从堆顶开始使用 down 对堆从上到下进行调整，最后调用 Pop 弹出最后一个元素，最后与 Push 一样，Pop 在定义时依然要保证弹出的时 Len - 1 位置的节点，并且修改堆长度，图 6 展示了 Pop 方法的执行过程。
// Pop removes and returns the minimum element (according to Less) from the heap.// The complexity is O(log n) where n = h.Len().// Pop is equivalent to Remove(h, 0).func Pop(h Interface) any {n := h.Len() - 1h.Swap(0, n)down(h, 0, n)return h.Pop()}
Remove && Fix
最后是 Remove 和 Fix 方法，Remove 方法中 L7 ~ L9 基本与 Fix 方法相同，所以我们放到一起来讲。先看 Remove 方法，它的作用与 Pop 方法相同，但是可以弹出指定位置的节点，首先找到最后一个节点的索引，如果目标节点是最后一个节点，那么直接 Pop，如果目标不是最后一个节点，那么与最后一个节点的值进行交换，接下来就是 Fix 的逻辑，首先在目标位置（ i ）进行 down 操作，而正如上文所述 down 返回一个 bool 值来判断 down 是否对堆进行了修改，如果 down 操作没有进行修改则说明目标位置在堆的最后一层，这时执行 up 操作向上层进行调整，到这里 Fix 操作的用途也就清楚了，选择合适的方法对堆进行调整来重新满足最小堆，图 7 展示了 Remove 时 i 在最后一层与中间一层的两种情况。
// Remove removes and returns the element at index i from the heap.// The complexity is O(log n) where n = h.Len().func Remove(h Interface, i int) any {n := h.Len() - 1if n != i {h.Swap(i, n)if !down(h, i, n) {up(h, i)}}return h.Pop()}// Fix re-establishes the heap ordering after the element at index i has changed its value.// Changing the value of the element at index i and then calling Fix is equivalent to,// but less expensive than, calling Remove(h, i) followed by a Push of the new value.// The complexity is O(log n) where n = h.Len().func Fix(h Interface, i int) {if !down(h, i, h.Len()) {up(h, i)}}
接下来我们来实际使用并测试一下这 5 个方法，下方为示例代码，首先按照接口要求定义一个类型然后绑定 5 个方法，sort 中的 Len、Less、Swap 以及上文提到的 Push 和 Pop，在 main 中我们定义一个 int 类型的 slice 让内容与上文图初始值相等，然后转类型为之前定义的 intHeap 类型，之后就是方法的使用了。
Testing
首先进行 Init 生成一个最小堆（结果为 L70），之后我们 Push 一个节点，值为 int 类型的 4（结果为 L72），然后进行 Pop 弹出堆顶元素（结果为 L74），使用 Remove 先弹出 index 1 再弹出 index 7，分别对于 down 和 up 的两种情况（结果为 L76～L77），输出结果都与上文图片相符，最后我们更改 index 4 和 index 7 来测试 Fix（结果为 L79~82）。
package mainimport ("container/heap""fmt")type intHeap []intfunc (h intHeap) Len() int {return len(h)}func (h intHeap) Less(x, y int) bool {return h[x] < h[y]}func (h intHeap) Swap(x, y int) {h[x], h[y] = h[y], h[x]}func (h *intHeap) Push(x any) {*h = append(*h, x.(int))}func (h *intHeap) Pop() any {old := *hn := len(old)x := old[n-1]*h = old[:n-1]return x}func main() {testing := intHeap([]int{8, 7, 5, 12, 4, 2, 11, 10, 1, 6, 3})fmt.Printf("Origin: %d \n", testing)heap.Init(&testing)fmt.Printf("Init minimum: %d \n", testing)fmt.Println()heap.Push(&testing, 4)fmt.Printf("Push 4: %d \n", testing)fmt.Println()heap.Pop(&testing)fmt.Printf("Pop: %d \n", testing)fmt.Println()heap.Remove(&testing, 1)fmt.Printf("Remove index 1: %d \n", testing)heap.Remove(&testing, 7)fmt.Printf("Remove index 7: %d \n", testing)fmt.Println()testing[4] = 1fmt.Printf("change index 4 = 1: %d \n", testing)heap.Fix(&testing, 4)fmt.Printf("Fix index 4: %d \n", testing)testing[7] = 3fmt.Printf("change index 7 = 3: %d \n", testing)heap.Fix(&testing, 7)fmt.Printf("Fix index 7: %d \n", testing)}// L37 --> Origin: [8 7 5 12 4 2 11 10 1 6 3]// L40 --> Init minimum: [1 3 2 7 4 5 11 10 12 6 8]//// L44 --> Push 4: [1 3 2 7 4 4 11 10 12 6 8 5]//// L48 --> Pop: [2 3 4 7 4 5 11 10 12 6 8]//// L52 --> Remove index 1: [2 4 4 7 6 5 11 10 12 8]// L55 --> Remove index 7: [2 4 4 7 6 5 11 8 12]//// L59 --> change index 4 = 1: [2 4 4 7 1 5 11 8 12]// L61 --> Fix index 4: [1 2 4 7 4 5 11 8 12]// L64 --> change index 7 = 3: [1 2 4 7 4 5 11 3 12]// L66 --> Fix index 7: [1 2 4 3 4 5 11 7 12]
至此 heap 包中的函数与方法已经全部阅读完毕，只要理解了 down 和 up 两个函数其它的内容很容易理解，不懂的地方只要去画图理解很容易就能理清逻辑关系，本周就到这里，我们下周再见。
若有收获，就点个赞吧
0 人点赞
