说到 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 的排序效果会更明显的变好一些。
回归到一般的业务工作中,我们构造待排序的数据集其实也非常花费时间,如果使用堆排序,可以省去一部分构造待排序数据集的时间开销。不过,还是需要合理评估,如果数据集本身特别小,只有几十个元素的排序,可能堆排序还会特别拉胯
参考文章:
- The Generic Way to Implement a Heap in Golang
- Usage of the Heap Data Structure in Go (Golang), with Examples