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包来构建和操作堆并可以实现堆排序等应用。