Go 堆数据结构使用

说到 container/heap 下的堆数据结构，让我们不需要从零开始实现这个数据结构。如果只是日常工作，其实还挺难用到堆的，更多的还是在写算法题的时候会用到。

基本概念

堆分为大顶堆和小顶堆，区分这两种类型方便我们处理问题。大顶堆的堆顶元素值最大，如果我们的业务场景是求TopN的最小值，我们就可以维系N个元素的大顶堆。这样的好处是，如果我们待排序的元素大于堆顶元素，直接忽略这个元素就可以。小顶堆也同理。

如图，小顶堆和大顶堆的直观展示。堆的结构是一颗完全二叉树，拿小顶堆来说，父节点要始终小于它的左右子节点，但左右子节点的大小是不明确的。

我们一般通过数组来存储完全二叉树，也就是用数组来存储堆结构，具体的表示关系，关键是看数组的 0 号元素是否要存储值。我们的实现都是数组的 0 号元素作为堆顶元素。用数组表示之后，具体的元素关系如下：

如果我们的存储是从下标为1开始的，数组的 1 号元素表示堆顶元素，对应的计算关系就会发生变化。左孩子节点为 2i，右孩子节点为 2i+1，父节点为 i/2。不过，go 语言数组 0 号位标识堆顶元素，和图式是一致的。

基础操作

堆的基本操作包括向堆中插入一个新元素，以及删除堆顶元素。熟悉了这两个操作之后，基本也就掌握了堆的实现。本质上，就是模拟了一个空元素，然后重新调整堆结构。插入是一个不断上浮的过程，删除是一个下沉的过程，注意看下面两个过程：

插入

我们向堆中插入新的元素 7，就是在完全二叉树的叶子结点上追加一个新的节点（追加在从左到右的最后），然后依次对这个新节点所在的子树做调整，使最小子树保持堆结构，最终到不需要调整结束，此时也就找到了元素 7 的最终位置。

从底层数组的角度来看的话，就是向数组中追加一个新的元素 7，然后基于父子节点关系，不断进行向上调整，指导找到 7 的最终位置。

删除堆顶

删除堆顶元素和插入类似，堆顶元素可以通过访问数组的第一个元素获取到，删除堆顶元素之后，堆顶元素就空了。堆的处理思路是将最后一个叶子节点放到堆顶元素，然后，不断进行所在子树的向下调整，最终所有子树都满足堆的特性。

如果所示，在删除堆顶元素 13 之后，我们将堆的末尾元素替换到堆顶的位置，然后，不断向下调整，最终找到元素 6 的合理位置。

go heap 实现

go 在 container/heap 做了官方的实现，如果要使用堆，只需要实现下面的接口，接口中匿名嵌套了排序接口，我们也需要将排序的三个接口声明实现一遍。

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

特别需要注意一点，在实现接口类型时，Push和Pop接受体的声明需要选择指针类型，因为这个过程需要对原始的值做扩展，在原始值得基础上做修改。

我们实现一个基本的堆排序接口，使用基础类型 int 来模拟堆排序。下面是代码部分，这部分是基本的堆实现，如果要使用堆的话，可以在这个基础上做扩展。

type IntHeap []int

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

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

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

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
}

实现堆的 Interface 接口之后，就是如何使用这个结构体了。我们简单来说明一下。下面这个例子，我们初始化了一个 IntHeap 数据类型，然后调用 heap.Init 方法做了堆排序，保证初始化的数组是按照小顶堆的顺序排序的。之后获取堆的长度，依次调用 heap.Pop 方法获取堆顶元素。

之所以每次都调用 Pop 方法，是因为在 Pop 完堆顶元素之后，剩余的堆元素还需要进行排序。小顶堆只知道堆顶元素最小，去掉堆顶元素之后，左右子树还是需要做调整的。

// output: 1,2,2,3,3,4,5,5,6,9,10,20,
func main() {
	h := &IntHeap{3, 2, 20, 5, 3, 1, 2, 5, 6, 9, 10, 4}
	heap.Init(h)

	length := h.Len()
	for i := 0; i < length; i++ {
		fmt.Printf("%d,", heap.Pop(h).(int))
	}
}

当然，我们也可以不使用 heap.Init 去做堆的初始化排序，我们可以向一个空元素的堆中依次插入元素来构建堆，效果其实是一样的。因为每次执行 heap.Push 就是在做堆排序。

func main() {
	nums := []int{3, 2, 20, 5, 3, 1, 2, 5, 6, 9, 10, 4}
	h := &IntHeap{}

	for _, val := range nums {
		heap.Push(h, val)
	}

	length := h.Len()
	for i := 0; i < length; i++ {
		fmt.Printf("%d,", heap.Pop(h).(int))
	}
}

性能比较

在 go 中如果不使用堆排序，我们要获取 TopN 个最大值元素，一般的做法是对数据集做好排序，然后，截取排序好的前 N 个元素。

如果使用堆排序，性能会更好吗？如果还以取 TopN 个最大值元素为例，我们可以构造一个 N 个元素的小顶堆，如果发现待排序的元素比堆顶元素小，可以直接舍弃。如果待排序的元素比堆顶元素大，执行 heap.Push 堆排序，完成之后并 heap.Pop 堆顶元素。

快排的方法，调用 sort.Slice 对整形数组直接排序，然后取 TopN

func FindTopNWithSort(nums []int, n int) []int {
	sort.Slice(nums, func(i, j int) bool {
		return nums[i] > nums[j]
	})

	result := make([]int, n)
	copy(result, nums[:n])
	return result
}

下面是堆排序方式，为了跟堆顶元素做比较，我们在 IntHeap 结构体上扩展新的方法 Top 来获取堆顶元素。这样能减少一些无效的堆排序。

func FindTopNWithHeap(nums []int, n int) []int {
	h := &IntHeap{}
	for _, val := range nums {
		// 根据堆顶元素做提前退出
		if h.Len() == n && h.Top().(int) > val {
			continue
		}

		heap.Push(h, val)
		if h.Len() > n {
			heap.Pop(h)
		}
	}

	return func() []int {
		initialLen := h.Len()
		result := make([]int, initialLen)
		for i := initialLen; i > 0; i-- {
			result[i-1] = heap.Pop(h).(int)
		}
		return result
	}()
}

我们首先验证这两个方法的正确性，确保可以返回 TopN 的最大值，然后在使用 benchmark 做性能测试。这里随机构造 600 个整数，然后用两种方法取 Top 30 的最大值

var nums []int
var result []int

func TestMain(m *testing.M) {
	maxVal := 6000
	rand.Seed(time.Now().Unix())

	nums = make([]int, maxVal)
	for i := 0; i < len(nums); i++ {
		nums[i] = rand.Intn(maxVal)
	}
	m.Run()
}

func BenchmarkFindTopNWithSort(b *testing.B) {
	k := 30
	nums2 := make([]int, len(nums))
	copy(nums2, nums[:len(nums)])
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		result = FindTopNWithSort(nums2, k)
	}
}

func BenchmarkFindTopNWithHeap(b *testing.B) {
	k := 30
	nums2 := make([]int, len(nums))
	copy(nums2, nums[:len(nums)])
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		result = FindTopNWithHeap(nums2, k)
	}
}

两种方式执行5次，可以发现，heap 的执行情况要稍微好一点。如果将待排序的数据集由600调整成6000，heap 的排序效果会更明显的变好一些。

回归到一般的业务工作中，我们构造待排序的数据集其实也非常花费时间，如果使用堆排序，可以省去一部分构造待排序数据集的时间开销。不过，还是需要合理评估，如果数据集本身特别小，只有几十个元素的排序，可能堆排序还会特别拉胯

参考文章：

1. The Generic Way to Implement a Heap in Golang
2. Usage of the Heap Data Structure in Go (Golang), with Examples