堆排序通常基于数组,其最优时间复杂度、平均时间复杂度、最坏时间复杂度均为 O(nlogn)

我们常说的优先队列就可以使用堆来实现

建堆

这里我们给定一个无序数组

arr := []int{9, 3, 1, 8, 2, 4, 0, 5, 6, 4, 7}

我们将在这个数组上建立一个小根堆,即最前面的元素最小

堆一般可等价一个完全二叉树,假设父节点在数组中的下标为 root,那么其两个孩子的下标分别为 root*2 + 1root*2 + 2

同时,根据一个给定的节点下标 x,可知其父节点的下标为 (x - 1) / 2 (这里是截去小数位整除)

如果一个数组的大小为 n,那么其倒数第二层的最后一个有孩子的节点的下标为 n/2 - 1

根据上述定义,可知一个堆用数组表示为:

heap := []Element{
    root, rootLeftChild, rootRightChild,
    rootLeftChildLeftChild, rootLeftChildRightChild ...}

我们将一个堆看做一个完全二叉树,那么第一步,我们从此树的倒数第二层开始构建一个小根堆

将问题分解,将给出根节点下标的子树进行下沉操作,将值较大的节点下沉到叶子节点

func down(arr []int, root, size int) {
	child := root*2 + 1
	for child < size {
		if child+1 < size && arr[child+1] < arr[child] {
			child++
		}
		if arr[child] > arr[root] {
			break
		}
		arr[root], arr[child] = arr[child], arr[root]
		root = child
		child = root*2 + 1
	}
}

arr 是原始未排序的堆,root 是子树的根节点在堆中的下标,size 是堆的大小,它用来限制节点下标的范围,构成了一个左闭右开的区间

我们需要保证下沉操作完成后,父节点小于其两个子节点,所以我们需要先找出子节点中较小者,并将其与父节点比较,若父节点大于当前子节点,交换,交换后两子节点中的较小者就会成为当前父子节点中的新父节点,它小于其他两个子节点(其中一个是被换下来的旧父节点)

if child+1 < size && arr[child+1] < arr[child] {
    child++
}

这之后判断当前父节点是否大于子节点中的较小者,若是,则交换两者,并将 root 指向当前最小子节点,child,指向其左子节点,递推往下,下沉较大值

基于上面的操作,我们可以将整个堆构建成一个最小堆,只需要从下往上挨个下沉较大的父节点(们)即可

func makeHeap(arr []int) {
	n := len(arr)
	for i := n/2 - 1; i >= 0; i-- {
		down(arr, i, n)
	}
}

在 Go 中,传递一个切片需要考虑修改其大小和容量会不会对后续操作产生影响,若有,最好传递其指针 &arr

经过上述操作,可得数组

[0 2 1 4 7 8 3 5 6 9]

可以看到,现在还并不是一个有序数列,但它已经是一个小根堆了!(如果再加点操作就可以用作一个优先队列了)

排序

现在的小根堆,看做一个完全二叉树,还不是一个有序数列的原因在于,子树与子树,子树的子树和子树的子树之间并没有联系

这是因为我们从下往上构建堆时并没有回溯的操作,两个子树(子堆)之间并没有形成一定的联系

我们通过如下算法将子树与子树,和子树的子树与子树的子树之间的联系打通:

从最后一个元素开始,将其与第一个元素 arr[0] 交换,然后从堆顶(树根)开始进行下沉操作,下沉到最下面的是最小的值

这一轮之后,对除了最后一个下沉完成的元素以外的其他元素进行如上操作

从上面可以分析出,我们将小根堆的堆顶(也就是最小值)全都扔在了数组最后,由此构建出来一个递减的数组

func HeapSort(arr []int) {
	n := len(arr)
	for i := n/2 - 1; i >= 0; i-- {
		down(arr, i, n)
	}

	for i := n - 1; i >= 1; i-- {
		arr[i], arr[0] = arr[0], arr[i]
		down(arr, 0, i-1)
	}
}

func down(arr []int, root, size int) {
	child := root*2 + 1
	for child < size {
		if child+1 < size && arr[child+1] < arr[child] {
			child++
		}
		if arr[child] > arr[root] {
			break
		}
		arr[root], arr[child] = arr[child], arr[root]
		root = child
		child = root*2 + 1
	}
}

后话

Go 的标准库中也提供了堆的模板,具体位于 container/heap 这个包中,我们可以很方便的实现一个优先队列

注意要使用后面重点标注的方法,不要使用自己实现的接口里的方法

type Heap []int

func (h Heap) Len() int {
	return len(h)
}

// 这里默认是下沉时取自孩子的操作,i 为右孩子,j 为左孩子,也就是,如果右孩子 i 小于左孩子 j
// 那么将我们例子中的孩子指针 child 指向 i
//
// 另外,如果改成大于,则会构建出一个大根堆
//
// 由于 std 添加了我们缺失的方法,所以更具体点,这会构建出一个优先队列
//
// 成就:发现 Go 的优先队列了!
func (h Heap) Less(i int, j int) bool {
	return h[i] < h[j]
}

func (h Heap) Swap(i int, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *Heap) Pop() any {
	res := (*h)[h.Len()-1]
	*h = (*h)[:h.Len()-1]
	return res
}

func (h *Heap) Push(x any) {
	*h = append(*h, x.(int))
}

需要注意的是,这里我们的 PushPop 方法只是简单的在底层类型 []int 上 append 和缩窄了,不能保证队列的有序性

为此我们需要用同一个包中的 heap.Push/heap.Pop 方法(当然你也可以直接在 Push 和 Pop 直接实现,但既然有能用的干嘛要自己写呢 :P)

e.g.

func main() {
    pq := &Heap{}
    heap.Push(pq, 1)
    fmt.Printf("%v\n", heap.Pop(pq))
}

这里有一点需要注意,如果是实现优先队列,只需要保证取出的值是当前集合中的最值(可能最大,也可能最小),所以不需要保证整个集合的有序性

所以上面的方法:

func (h Heap) Less(i int, j int) bool {
	return h[i] < h[j]
}

会构建出一个递增的优先队列,而放到堆排序中,会构建出一个递减的数组(因为 Go 的 container/heap 包并不会进行排序,若用作优先队列则不会出现将最小值全部放到底层数组末尾的操作)

所以如果想保证取出的优先队列能够保证有序性,需要使用额外的排序方法,或者逐次 heap.Pop 获取所有元素

sort.Ints()/slices.Sort()