实现一个友好的堆
source link: https://colobu.com/2024/03/07/implement-a-heap-in-go/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
实现一个友好的堆
在上一篇文章中，我吐槽了Go标准库的堆实现，基于“you can you up, no can no BB”的理论，这篇文章我来实现一个友好的堆。
我们使用堆的时候，一般希望有Heap这样一个对象，并且能指定它是“小根堆”或者"大根堆"。我们希望这个类型有Push和Pop方法，可以加入一个元素或者弹出(最小的)元素。
我们期望这个Heap支持泛型的，任何可以比较的类型都可以使用。
处于简化的考虑，我们实现的Heap不考虑线程安全。如果要保证线程安全，可以使用sync.Mutex来保护Heap的操作。
我们实现的Heap类型的操作基于标准库的操作，只不过我们封装了一下，让它更加友好。
我们能够基于既有的一个slice创建Heap，也可以基于一个空的Heap创建一个新的Heap。
最终我们实现了一个友好的堆，你可以在github上查看它的代码binheap。
首先定义一个binHeap,这是一个泛型的slice,用来保存堆的元素，这样用户就不用定义这样一个类型了，简化了用户的使用。默认是小根堆。所有元素类型需要满足cmp.Ordered接口，可以进行大小比较。这个接口是标准库中的接口，如果你还不知道cmp包，那么需要刷新刷新Go新的变化了：
package heapimport ( "cmp"type binHeap[V cmp.Ordered] []Vfunc (h binHeap[V]) Len() int { return len(h) }func (h binHeap[V]) Less(i, j int) bool { return h[i] < h[j] }func (h binHeap[V]) Swap(i, j int) { h[i], h[j] = h[j], h[i] }func (h *binHeap[V]) Push(x V) { *h = append(*h, x)func (h *binHeap[V]) Pop() V { old := *h n := len(old) x := old[n-1] *h = old[0 : n-1] return x这样我们就可以定义一个BinHeap类型，它是binHeap的封装，它有maxHeap字段，用来表示是小根堆还是大根堆。BinHeap类型有Push和Pop方法，可以加入一个元素或者弹出(最小的)元素。
type BinHeap[V cmp.Ordered] struct { maxHeap bool binHeap binHeap[V]// NewBinHeap returns a new binary heap.func NewBinHeap[V cmp.Ordered](opts ...BinHeapOption[V]) *BinHeap[V] { h := &BinHeap[V]{} for _, opt := range opts { opt(h) return h// Len returns the number of elements in the heap.func (h *BinHeap[V]) Push(x V) { h.binHeap.Push(x) sift_up[V](&h.binHeap, h.binHeap.Len()-1, h.maxHeap)// Push pushes the element x onto the heap.func (h *BinHeap[V]) Pop() V { n := h.binHeap.Len() - 1 h.binHeap.Swap(0, n) sift_down[V](&h.binHeap, 0, n, h.maxHeap) return h.binHeap.Pop()另外还附送了两个常用的方法Len和Peek，Len返回堆的大小，Peek返回堆顶元素但是并不会从堆中移除它，在和堆的最小值做比较的时候很有用。
// Len returns the number of elements in the heap.func (h *BinHeap[V]) Len() int { return h.binHeap.Len()// Peek returns the element at the top of the heap without removing it.func (h *BinHeap[V]) Peek() (V, bool) { var v V if h.Len() == 0 { return v, false return h.binHeap[0], true最后，我们还提供了一个BinHeapOption类型，用来设置BinHeap的属性，比如是小根堆还是大根堆；为了提高性能，如果预先已经知道堆的大小，可以在初始化的时候就进行设置。
// WithMaxHeap returns a BinHeapOption that configures a binary heap to be a max heap.func WithMaxHeap[V cmp.Ordered](h *BinHeap[V]) *BinHeap[V] { h.maxHeap = true return h// WithMinHeap returns a BinHeapOption that configures a binary heap to be a min heap.func WithCapacity[V cmp.Ordered](n int) BinHeapOption[V] { return func(h *BinHeap[V]) *BinHeap[V] { if h.binHeap == nil { h.binHeap = make(binHeap[V], 0, n) return h这样我们就实现了一个友好的堆。
当然，如果你已经有了一个slice: []V, 想把它转换成堆，并且在这个slice上进行堆操作，那么你可以使用下面的方法：
// NewBinHeapWithInitial returns a new binary heap with the given initial slice.func NewBinHeapWithInitial[V cmp.Ordered](s []V, opts ...BinHeapOption[V]) *BinHeap[V] { h := &BinHeap[V]{} h.binHeap = binHeap[V](s) for _, opt := range opts { opt(h) n := len(s) for i := n/2 - 1; i >= 0; i-- { sift_down[V](&h.binHeap, i, n, h.maxHeap) return h堆的操作sift_down和sift_up的堆和核心操作，也来源子标准库的代码，只不过我把它们改成成泛型的函数了：
func sift_up[V cmp.Ordered](h *binHeap[V], j int, maxHeap bool) { less := h.Less if maxHeap { less = func(i, j int) bool { return !h.Less(i, j) } i := (j - 1) / 2 // parent if i == j || !less(j, i) { break h.Swap(i, j)func sift_down[V cmp.Ordered](h *binHeap[V], i0, n int, maxHeap bool) bool { less := h.Less if maxHeap { less = func(i, j int) bool { return !h.Less(i, j) } i := i0 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 && less(j2, j1) { j = j2 // = 2*i + 2 // right child if !less(j, i) { break h.Swap(i, j) return i > i0Recommend
About Joyk
Aggregate valuable and interesting links.
Joyk means Joy of geeK