Golang heap堆源码解析

堆作用

Golang 的堆是用二叉树来实现，heap包实现了保证根节点为最大/最小值即大根堆/小根堆，而后续其他的节点顺序不能保证，要想严格排序，需要自行实现额外逻辑。

使用场景

最为经典的场景就是寻找最大或最小值

用法

实现 heap.Interface接口
必须使用heap.Init初始化
必须通过heap.Push和heap.Pop接口来添加删除元素

 var _ heap.Interface = (*hp)(nil)type hp struct {sort.IntSlicedata []int
}func (h *hp) Push(x any) {h.IntSlice = append(h.IntSlice, x.(int))
}func (h *hp) Pop() any {h.IntSlice = h.IntSlice[:len(h.IntSlice)-1]return h.IntSlice[0]
}func (h *hp) Less(i, j int) bool {return h.data[h.IntSlice[i]] > h.data[h.IntSlice[j]]
}func maxSlidingWindow(nums []int, k int) []int {h := &hp{make(sort.IntSlice, k),nums,}for i := 0; i < k; i++ {h.IntSlice[i] = i}heap.Init(h)r := make([]int, 0, len(nums)-k+1)r = append(r, h.data[h.IntSlice[0]])for j := k; j < len(nums); j++ {heap.Push(h, j)for h.IntSlice[0] <= j-k {heap.Pop(h)}r = append(r, h.data[h.IntSlice[0]])}return r
}

核心源码

// Interface 堆的底层数据结构为二叉树
type Interface interface {sort.InterfacePush(x any) // 添加元素Pop() any   // 移除并返回元素
}// Init 方法，初始化后保证父节点和子节点是有序的（不同层数据是有序的）
func Init(h Interface) {// heapifyn := h.Len()for i := n/2 - 1; i >= 0; i-- { // 循环的数据为 除了叶子节点以外所有的节点（保证所有节点都是有序）down(h, i, n)}
}// 插入元素，并把当前元素上浮到合适位置
func Push(h Interface, x any) {h.Push(x)up(h, h.Len()-1)
}// Pop 移除根节点
// 把根节点移到尾部，并把尾部节点下沉到合适位置
func Pop(h Interface) any {n := h.Len() - 1h.Swap(0, n)down(h, 0, n)return h.Pop()
}// Remove 把待删除节点和尾部节点交换，并把尾部节点上浮或者下沉到合适位置
func Remove(h Interface, i int) any {n := h.Len() - 1if n != i {h.Swap(i, n)if !down(h, i, n) {up(h, i)}}return h.Pop()
}// Fix 对于任意节点重新寻找位置（下沉、上浮）
func Fix(h Interface, i int) {if !down(h, i, h.Len()) {up(h, i)}
}// up
//
//	@Description: j 节点 上浮
//	@param h: 堆
//	@param j: 待操作节点
func up(h Interface, j int) {for {i := (j - 1) / 2 // 父节点if i == j || !h.Less(j, i) {break}h.Swap(i, j)j = i}
}// down 方法
//
//	@Description: 当前节点i0 根据大小关系是否下沉（直到叶子节点）
//	@param h: 堆
//	@param i0: 当前待处理节点
//	@param n：节点数量
//	@return bool: 是否进行了下沉
func down(h Interface, i0, n int) bool {i := i0for {j1 := 2*i + 1          // i 的左子节点if j1 >= n || j1 < 0 { // j1 < 0 after int overflow  //保证左子节点存在/有效break}j := j1                                     // left childif j2 := j1 + 1; j2 < n && h.Less(j2, j1) { // j2 为i的右子节点j = j2 // = 2*i + 2  // right child}if !h.Less(j, i) {break}h.Swap(i, j)i = j}return i > i0
}