Go语言里面的Heap实现

[bruce-16]

Heap

最大（小）堆：通俗的就是说，父节点比左右子节点都大或者小的完全二叉树。

完全二叉树（特点）：

1. 只允许最后一层有空缺结点且空缺在右边，即叶子结点只能在层次最大的两层上出现；
2. 1. 对任一结点，如果其右子树的深度为j，则其左子树的深度必为j或j+1。 即度为1的点只有1个或0个

最近学习图的单源最短路径算法（我看的是Dijkstra算法）的时候，有看到讲解算法的博主在里面使用最小堆来记录每个顶点距离原点的最小距离，然后每次从最小堆顶中取出最小距离的点，再对邻接表进行循环。因为Java的提供优先队列的实现是没有更新的接口，所有博主自己写了一套，Go语言的标准库里面也是有Heap的实现的，提供了相应的方法。Go语言的实现只是引子，重点还是思路。

最小堆或者最大堆的核心操作

不管是插入、删除和更新操作，最重要的就是维持最小堆或者最大堆的特性（每个节点都比左右子节点小或者大）。所以先看看，当一个节点被操作后，违反了这个特性，那么就接下来的步骤就是从该节点开始调整整棵树了。比如有这样的一棵树：

假设图中被红框标注的那个节点刚被修改了，这个节点的值相比父节点较小，所以已经不满足了最小堆的定义了，就需要调整了。调整的分为两类：1、上浮； 2、下沉。像上面这种情况，就考虑上浮。

上浮的步骤

1. 拿变更的接节点与父元素进行比较，如果比父元素小（大顶堆就是大）的话就与父元素交换位置，否则不交换。
2. 如果1步骤发生交换，那么将父元素作为比较点，去循环第一步骤（往上），一直到某个顶点没有父元素（根元素）。如果1步骤没有发生交换，则结束。

---

如下图，但把根节点删除的时候，就需要把最后一个节点替换到根节点，然后再对根节点进行下沉，以维持最小（大）堆的平衡。

下浮的步骤

1. 将当前节点与左右子节点中最小的子节点比较（注意边界，有可能没有右节点），如果子节点有比当前根节点小的值，那么就进行交换。否则不交换。
2. 如果1步骤发生交换，则定位到交换后的位置（子节点），再对该节点循环1步骤直到没有子节点（叶子节点）。如果1步骤没有交换则结束。

Heap的源码

将指定类型初始化为最小（大）顶堆

type Interface interface {
	sort.Interface
	Push(x interface{})
	Pop() interface{}
}
func Init(h Interface) {
	n := h.Len()
	for i := n/2 - 1; i >= 0; i-- {
		down(h, i, n) // 对节点i进行下沉操作
	}
}

Init函数接收一个类型为Interface的形参（h），我们先看一下Interface的结构它是由sort.interface和Push、Pop函数组成，也就是这个实现这个接口的类型需要有：Len()、Less(i, j int) bool、Swap(i, j int)、Push(x interface{})和Pop() interface{}这几个方法。

现在我们继续说Init函数，该函数很简单，就一个循环，但是起点比较奇怪，然后再遍历到的节点进行下沉（down函数）操作。首先看为什么是从n/2-1这个节点开始。

对于一个完全二叉树，如果节点个数为n那么它的所以叶子节点就是n/2到n包含两者，这些节点都是叶子节点。那么n/2-1也就是从后往前第一个有子节点的节点。叶子节点没有子节点，不需要做下沉操作，所以循环的开始应该是从后往前第一个有子节点的节点开始进行下沉操作。

添加元素

func Push(h Interface, x interface{}) {
	h.Push(x)
	up(h, h.Len()-1)  // 上浮操作
}

将元素添加到尾部，这里是调用的类型为Interface变量h的Push方法，也就是使用的时候需要手动去实现的，然后将最后这个元素（最后一个子节点）进行上浮操作。

删除元素

func Pop(h Interface) interface{} {
	n := h.Len() - 1
	h.Swap(0, n) // 将第一个元素与最后一个元素交换
	down(h, 0, n) // 对顶部进行下沉操作
	return h.Pop() // 返回并删除最后一个元素，也就是被交换下来最开始要删除的顶部元素
}

这里将第一个元素与最后一个元素进行交换也就是把顶部的节点和最后一个子节点进行了交换，然后再对交换后的顶部节点进行下沉操作来维护最小顶堆的平衡。然后再把最后一个叶子节点删除并返回，也就是被交换下来的之前的顶部节点。

移除某一个节点

func Remove(h Interface, i int) interface{} {
	n := h.Len() - 1
	if n != i {
		h.Swap(i, n) // 待删除的节点，与最后一个节点进行交换
		if !down(h, i, n) {
			up(h, i)
		}
	}
	return h.Pop()
}

到这里可以看到，只要删除节点，就先把这个节点与最后一个节点进行交换，然后在对这个节点进行调整。但是这里地方有点特殊，可能需要下沉操作后还需要上浮操作。因为删除的这个节点并不一定是顶部节点或者叶子节点。为什么需要有!down(h, i, n) 这个判断，因为如果一个节点进行了下沉操作，则说明该节点比子节点还要大，由于其他节点之前都是按照最小顶堆放置的，也就是说，如果该节点比子节点还要大，那么它就肯定比自己的父节点还要大。所以这里就做了优化，如果这个节点进行了下沉操作，则不需要进行上浮了，否则再进行上浮。其实判断是否有过上浮操作再判断是否需要下沉是一样的。

调整某个节点的特性

func Fix(h Interface, i int) {
	if !down(h, i, h.Len()) {
		up(h, i)
	}
}

与删除差不多，就是对该节点进行下沉或者上浮。

下沉操作的实现

func down(h Interface, i0, n int) bool {
	i := i0
	for {
		j1 := 2*i + 1
		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
			break
		}
		j := j1 // left child
		if 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
}

首先有一个循环，循环的结束条件为节点没有子节点了或者节点经过比较不需要和子节点进行交换。现在看一下逻辑:

1. 首先找到左子节点的下标j1 := 2*i + 1
2. 判断是否有右几点，并且是否比左节点要小j2 := j1 + 1; j2 < n && h.Less(j2, j1)，如果判断成立，则节点和右子节点进行比较，否则和左节点进行比较j = j2 // = 2*i + 2 // right child。
3. 和子节点比较，如果符合条件，则进行交换if !h.Less(j, i) { break } h.Swap(i, j)
4. 如果有交换则将节点更新为被交换的子节点下标，循环1-3步骤。

上浮操作实现

相对于下沉操作，上浮的操作相对简单，因为只需要比较该节点与父节点的值，父节点只有一个。

func up(h Interface, j int) {
	for {
		i := (j - 1) / 2 // parent
		if i == j || !h.Less(j, i) {
			break
		}
		h.Swap(i, j)
		j = i
	}
}

循环的结束条件需要注意一下，当节点为根的时候，也就是j为0的时候，(0-1)/2是为0的，也就是最后计算出来的i也是为0的，所有 i == j成立，循环结束。所有循环结束的条件是，当前节点没有父节点了（根节点）或者不需要和父节点进行交换（节点的值大于父节点的值）。

使用Heap

既然写到这里了，那就看看go里面怎么使用它吧，我对go不是很熟，只是用来练习一些数据结构和算法而已，下面的一些代码并不是最佳实践。

先定义一个类型，实现Interface需要的几个方法

import "container/heap"

// PriorityQueue 优先队列
type PriorityQueue []int

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

func (pq PriorityQueue) Less(i, j int) bool {
	return pq[i] < pq[j]
}

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

// Pop 删除最后一个元素
func (pq *PriorityQueue) Pop() interface{} {
	n := pq.Len()
	item := (*pq)[n-1]
	*pq = (*pq)[0 : n-1]
	return item
}

// Push 在最后添加元素
func (pq *PriorityQueue) Push(v interface{}) {
	*pq = append(*pq, v.(int))
}

这里定义了一个类型PriorityQueue，本质是一个int型切片，然后实现了需要实现的几个方法，这样就可以使用heap提供的方了。

使用heap的方法

import "container/heap"

// PriorityQueue 优先队列
type PriorityQueue []int
...


func main() {
    queue := []int{4, 13, 14, 12, 9, 11, 8, 7, 10}
    // 堆化
    heap.Init(queue)
    // 添加元素
    heap.Push(queue, 5)
   // 顶部出堆
    heap.Pop(queue)
   // 删除某个节点
   heap.Remove(queue, 5)
   ...
}

总结

就这样吧，重点是上浮和下沉两个操作。有了最小（大）堆，还可以利用它进行排序（堆排序），也就是每次Pop顶元素就行。

下一篇准备写求图的单源最短路径的经典算法：Dijkstra。

blog地址