Go数据结构详解之“堆（Heap）” ｜ 豆包MarsCode AI 刷题

概念

什么是堆？

堆是一种特殊的完全二叉树，堆可以进一步细分为大顶堆和小顶堆

• 大顶堆的特点是，父节点的值总是大于或等于其子节点的值。
• 小顶堆的特点是，父节点的值总是小于或等于其子节点的值。

小顶堆的首个元素是整个堆中的最小值，大顶堆的首个元素是整个堆的最大值，利用这种特性，能够快速的访问和提取一组数据中的最小或最大值的元素。

通常我们使用数组来实现堆，你们可以通过下面这些规律直接计算出当前节点的父节点和子节点的索引：

若当前节点的下标为i存在

• 父节点的索引为：(i - 1) / 2
• 左子节点索引为： i * 2 + 1
• 右子节点索引为： i * 2 + 2

堆的两大操作——push和pop

**push**

1. 先将元素添加到堆的末尾
2. 将元素与父节点比较
3. 假设是小顶堆，如果元素比父节点小，则元素与父节点交换；假设是大顶堆，如果元素比父节点大，则元素与父节点交换
4. 重复此过程，直到元素达到根节点或元素满足对属性

pop

1. 先将根节点与堆的最后一个元素交换
2. 再删除最后一个元素
3. 对新的根节点执行"向下堆化"，维持整个堆的特性。

ps：一些常见的二叉树

满二叉树

二又树的所有叶子节点都在最后一层，并且结点总数= 2^n-1（n为层数），则我们称为满二叉树

完全二叉树

如果该二叉树的所有叶子节点都在最后一层或者倒数第二层，而且最后一层的叶子节点在左边连续，倒数第二层的叶子节点在右边连续，我们称为完全二叉树。

我们可以看出满二叉树是特殊的完全二叉树

平衡二叉树也叫平衡二叉搜索树(Self-balancing binary searchtree)又被称为 AVL 树，可以保证查询效率较高。

具有以下特点:它是一 棵空树或它的左右两个子树的高度差的绝对值不超过1，并且左右两个子树都是一棵平衡二叉树。平衡二叉树的常用实现方法有红黑树、AVL、替罪羊树、Treap、伸展树等。

Go语言中的堆

堆可以用数组、链表、二叉搜索树、平衡二叉搜索树等数据结构来实现。但通常选用数组来实现堆，因为可以通过代数方式直接计算出当前节点的父节点、左孩子和右孩子的位置，而不是通过遍历查找父节点与孩子节点的位置。因此它在时间和空间上都是很高效的。

Go语言中堆的实现，可以自己直接实现，也可以利用container/heap包实现。container/heap包提供了堆操作的接口和方法

接下来对container/heap包的源码分析

Go中heap的接口定义为

// The Interface type describes the requirements
// for a type using the routines in this package.
// Any type that implements it may be used as a
// min-heap with the following invariants (established after
// Init has been called or if the data is empty or sorted):
//
//	!h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
//
// Note that Push and Pop in this interface are for package heap's
// implementation to call. To add and remove things from the heap,
// use heap.Push and heap.Pop.
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)
}

对Go自定义排序比较熟悉对上面的接口并不陌生，heap中嵌入这个主要是因为在进行heapify(翻译为"堆化")时，需要使用Less()方法进行节点值的比较，使用Swap()方法进行节点位置之间的交换。

// Init establishes the heap invariants required by the other routines in this package.
// Init is idempotent with respect to the heap invariants
// and may be called whenever the heap invariants may have been invalidated.
// The complexity is O(n) where n = h.Len().
func Init(h Interface) {
	// heapify
	n := h.Len()
    // i从最后一个父节点n/2 - 1开始，因为((n/2) - 1) * 2 + 2 = n
	for i := n/2 - 1; i >= 0; i-- {
        // 从 n/2 - 1 到 0 开始进行堆化
		down(h, i, n)
	}
}
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
}

堆化的过程中，我们可以看到传入down()的有堆对象、当前节点和堆的大小。先判断当前节点的左子节点是否满足边界条件，如何判断当前节点的左右节点（如果存在右节点）的大小，选择出最大\最小和当前节点判断大小，如果当前节点小于\大于就交换他们。不断传入知道到达堆边界。这个还是好理解的，其实就是我们构建大小顶堆的过程。

Push

// 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)
}

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
	}
}

接下来我们来看看堆的push操作的逻辑，首先我们先调用h.Push()这个是我们自己去实现的添加方法，然后进入up()进行构建。传入最后一个节点，通过该节点找到该节点的父节点，比较该节点和父节点的大小来决定否进行交换，如果满足不交换的条件（看你是大顶堆还是小顶堆）就推出循环，否则会一直循环求出父节点然后比较，直到到达根节点或者满足不交换条件。

Pop

// 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()
}

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
}

Pop也比较好理解，将根节点和最后一个节点进行交换，然后剩余部分进行向下重建，这个就同理所得，不再赘述了。

我们还会注意到除了常见的Pop和Push，container/heap还提供了Remove和Fix

---

Remove 和 Fix

Remove方法用于从堆中移除并返回索引为i的元素。

Fix方法用于在堆中某个元素的值改变后，重新建立堆的顺序。这通常比先移除再添加新值更高效

// Remove removes and returns the element at index i from the heap.
// The complexity is O(log n) where n = h.Len().
func Remove(h Interface, i int) any {
	n := h.Len() - 1
	if n != i {
		h.Swap(i, n)
		if !down(h, i, n) {
			up(h, i)
		}
	}
	return h.Pop()
}

// Fix re-establishes the heap ordering after the element at index i has changed its value.
// Changing the value of the element at index i and then calling Fix is equivalent to,
// but less expensive than, calling Remove(h, i) followed by a Push of the new value.
// The complexity is O(log n) where n = h.Len().
func Fix(h Interface, i int) {
	if !down(h, i, h.Len()) {
		up(h, i)
	}
}

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
}

在里面的关键的、相同的一步就是先down()后up()，也就是说我们操作的i并非是节点，在与最后一个节点进行交换后既有可能是向下操作也可能是向上操作，调用down()方法尝试将新的h[i]下沉到正确的位置。如果down()没有改变h[i]的位置（即h[i]已经是正确的位置），则可能需要调用up()方法将其上浮。但在这个实现中，如果down()没有移动元素，则不调用up()，因为此时h[i]已经正确。

ps：从官方注释中我们也能知道Pop和Push的时间复杂度均为O(log<sub>2</sub> n)

堆的实际应用场景

+ 优先级队列： 例如任务调度器，即按任务优先级大小进行调度的任务调度器。告警系统中，按消息重要性优先发送重要信息等。 + 延时队列: 延迟队列是利用优先级队列，根据元素中的时间戳的先后顺序进行弹出。 + 堆排序: 利用堆排序，时间复杂度是O(n log₂ n)

优先队列

1. **堆作为优先级队列的实现**： - 堆是优先级队列的一种高效实现方式。在优先级队列中，查找和删除具有最高（或最低）优先权的元素的操作，可以通过堆的堆顶元素来高效地完成。 - 在堆中，插入一个新元素和删除堆顶元素的操作都可以通过对堆的调整（堆化）来维护堆的性质，从而实现优先级队列的功能。 2. **优先级队列的底层数据结构**： - 在许多编程语言中，优先级队列的底层实现都采用了堆这种数据结构。例如，在Java中，可以使用`PriorityQueue`类来创建优先级队列，而其底层实现就是基于堆的。

优先级队列和堆在多个领域都有广泛的应用。例如，在任务调度、Huffman编码、堆排序等算法中，优先级队列和堆都发挥着重要的作用。

package main

import (
	"container/heap"
	"fmt"
)

type Item struct {
	priority    int
	name string
}

// 优先级队列底层用数组实现
type PriorityQueue []*Item

// 先实现heap.Interfach的五个接口Len(),Less(),Swap(),Push(),Pop()
func (pq PriorityQueue) Len() int { return len(pq) }
func (pq PriorityQueue) Less(i, j int) bool {
  // 优先级值(priority)越大越优先(大顶堆), 也就是完全二叉树中，上层的节点的优先级大于下层节点的优先级值
	return pq[i].priority > pq[j].priority
}
func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
}

func (pq *PriorityQueue) Push(x interface{}) {
	Item := x.(*Item)
	*pq = append(*pq, Item)
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	Item := old[n-1]
	*pq = old[0 : n-1]
	return Item
}

func main() {
	pq := make(PriorityQueue, 0)
	heap.Init(&pq)

	// 插入元素
	Items := []Item{
		{1, "low"},
		{3, "high"},
		{2, "medium"},
	}
	for i := range Items {
		heap.Push(&pq, &Items[i])
	}

	// 按优先级从大到小的弹出元素
	for pq.Len() > 0 {
		Item := heap.Pop(&pq).(*Item)
		fmt.Printf("Item: %s\n", Item.name)
	}
}

-----------输出----------
Item: 高优先级
Item: 中优先级
Item: 低优先级

堆排序

这个就好理解了，先将需要pop的数字（根节点）交换到数组最后一位，然后重建前面数组。我们借助LeetCode的一道题来求解一下试试。

import (
	"container/heap"
	"sort"
)

type hp struct {
	sort.IntSlice
}

func (h *hp) Less(i, j int) bool {
    return h.IntSlice[i] > h.IntSlice[j]
}

func (h *hp) Push(v interface{}) {
	h.IntSlice = append(h.IntSlice, v.(int))
}

func (h *hp) Pop() interface{} {
	a := h.IntSlice
	v := a[len(a)-1]
	// Pop
	h.IntSlice = a[:len(a)-1]
	return v
}

func findKthLargest(nums []int, k int) int {
	q := &hp{nums}
	heap.Init(q)
	var ans any
	for i := 0;i < k;i++ {
		ans = heap.Pop(q)
	}
	return ans.(int)
}