golang实现大顶堆只看这篇文章就够了

前言

通过本篇文章，你将学会：

1. 初始化大顶堆
2. 弹出堆顶元素
3. 往堆中插入元素
4. 堆排序

学习的前提是你已经知道在构建好的堆中调整单个元素的原理，也就是下沉(down)操作和上浮(up)操作。

正文

在"container/heap"标准库中有实现堆的Init,Pop,Push,Fix等方法，只需要向heap的Init,Pop,Push,Fix等方法中传入实现了Interface的结构类型就可以使用heap的方法来初始化和调整堆

heap.Interface类型如下

type Interface interface {
	sort.Interface
	Push(x any) // add x as element Len()
	Pop() any   // remove and return element Len() - 1.
}

sort.Interface类型如下

type Interface interface {
	Len() int
	Less(i, j int) bool
	Swap(i, j int)
}

所以我们的自定义类型需要实现上面五个方法

type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h IntHeap) Less(i, j int) bool { return h[i] > h[j] } // 大顶堆
//func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] } // 小顶堆

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
}

我们一个一个来解释type IntHeap []int表明我们底层使用数组来存储堆元素
len()返回IntHeap 的长度
Swap(i, j int)交换IntHeap 中下标为i和j的元素
Less(i, j int)如果下标i的元素“小于”下标j的元素，则返回true，否则返回false。返回true则表明i不会在j的下方。通过控制Less函数的返回可以实现大顶堆或者小顶堆
Push(x any)函数往IntHeap的末尾插入元素
Pop()函数取出IntHeap末尾的元素

Pop之所以这样实现，是因为在heap包的源码中，Pop在调用IntHeap的Pop之前进行了一些操作：
heap.Push函数会往堆尾插入元素，如何对这个元素进行上浮(up)操作
heap.Pop函数会先交换堆首和堆尾元素，然后再对堆首元素进行下沉(down)操作，所以我们的IntHeap类型是往堆尾插入的。

// 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 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() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()
}

up和down代码如下

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
	}
}

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
}

golang实现堆的代码

package main

import (
	"container/heap"
	"fmt"
)

type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h IntHeap) Less(i, j int) bool { return h[i] > h[j] } // 大顶堆
//func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] } // 小顶堆

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{6, 1, 5, 3, 8, 7, 2}
	heap.Init(h)
	heap.Push(h, 4)
	for h.Len() > 0 {
		fmt.Printf("%d ", heap.Pop(h))
	}
}

输出：

8 7 6 5 4 3 2 1

堆排序

只要实现了下沉操作，就可以通过对非叶子节点进行下沉来初始化堆，通过将堆首元素弹出并置于堆尾即可实现堆排序

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

func down(values []int, n, i int) {
	largest := i
	l := 2*i + 1
	r := 2*i + 2
	if l < n && values[l] > values[largest] {
		largest = l
	}
	if r < n && values[r] > values[largest] {
		largest = r
	}
	if largest != i {
		values[i], values[largest] = values[largest], values[i]
		down(values, n, largest)
	}
}

也可以使用heap包来实现

func HeapSort2(values []int) {
	h := IntHeap(values)
	heap.Init(&h)
	n := h.Len()
	for i := 0; i < n; i++ {
		heap.Pop(&h)
	}
}

测试HeapSort2

func main() {
	var nums = []int{6, 1, 5, 3, 8, 7, 2, 4}
	HeapSort2(nums)
	fmt.Println(nums)
}

输出：

[1 2 3 4 5 6 7 8]

总结

通过这篇文章，可以学会如何使用heap包来构建和操作堆，并可以实现堆排序等应用。