Go语言中的堆实现与应用

Shawn摘要

由于Go没有直接提供像 Java 的 PriorityQueue 或 C++ 的 std::priority_queue 这样的封装好的堆数据结构，需要自己手动实现一些东西

Go 提供了 container/heap 包，允许开发者根据具体需求自定义堆的行为。

以下是我写的215. 数组中的第K个最大元素，使用 container/heap 包来实现堆的效果

type hp struct {
	sort.IntSlice
}

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

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

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

func findKthLargest(nums []int, k int) int {
	/*
	这里使用字面量初始化相当于使用 nums 初始化 IntSlice 匿名字段
	This composite literal allocates a new struct instance with the given values.
	You may either specify the field values in order or use their names.
	 */
	h := hp{nums}
	heap.Init(&h)
	for i := 0; i < k-1; i++ {
		heap.Pop(&h)
	}
	return heap.Pop(&h).(int)
}

匿名字段

这里的IntSlice是一个anonymous field，也就是匿名字段，如果我们写的是

1
2
3

type hp struct {
    nums sort.IntSlice
}

这样就不是继承，也就没有Len, Less, Swap方法

匿名字段相当于结构体里有一个和类型同名的字段也就是IntSlice，下面这样写也是ok的，但也不算继承，同样没有Len, Less, Swap方法

1
2
3

type hp struct {
	IntSlice sort.IntSlice
}

也就是说，Go语言中只有结构体里单放一个类型才叫继承

定义堆还可以是以下这种方式

type hp []int

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

func (h *hp) Pop() any {
	v := (*h)[len(*h)-1]
	*h = (*h)[:len(*h)-1]
	return v
}

func (h *hp) Len() int {
	return len(*h)
}

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

func (h *hp) Swap(i, j int) {
	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
}

func findKthLargest(nums []int, k int) int {
	// 不能使用 hp{nums}, 一定要用应该传递类似 &hp{2, 1, 5} 这种形式, {} 里面应该是 int.
	// 将 nums 通过 hp(nums) 转换成hp, 让同样的数据 nums 有了Push, Pop等方法
	h := hp(nums)
	heap.Init(&h)
	for i := 0; i < k-1; i++ {
		heap.Pop(&h)
	}
	return heap.Pop(&h).(int)
}

Go中不建议同时使用value receiver和pointer receiver

这里的heap.Init, heap.Pop中必须要传递&h，因为我们写的方法Len, Less, Swap, Push, Pop都是写的pointer receiver，而Init, Pop底层会调用我们写的这些方法，所以应该传递指针

同时，Push和Pop方法也只能传递指针

不然会提示 Assignment to the method receiver doesn't propagate to other calls

在 Go 语言中，方法接收者有两种形式：指针接收者（pointer receiver）和值接收者（value receiver）

指针接收者

修改接收者的内部状态：
- 当一个方法通过指针接收者调用时，它可以直接修改接收者的内部状态。
- 如果使用值接收者，方法接收的是接收者的一个副本，任何对副本的修改都不会影响原始接收者。
效率：
- 对于大型结构体，使用指针接收者可以避免每次调用方法时都复制整个结构体，从而提高效率。

结合代码解释

指针接收者的使用：
- Push 方法使用指针接收者 *hp，因此它可以直接修改 hp 的内部切片。
- Pop 方法同样使用指针接收者 *hp，以便直接修改切片的长度。
如果使用值接收者：
- 如果将 Push 和 Pop 方法的接收者改为值接收者 h hp，则这些方法将对 hp 的副本进行操作，而不是原始对象。
- 任何对副本的修改都不会影响调用这些方法时传递的原始 hp 对象，从而导致修改无法生效。

值接收者的问题

假设我们将 Push 和 Pop 方法改为值接收者，代码如下：

func (h hp) Push(v any) {
	h = append(h, v.(int))
}

func (h hp) Pop() any {
	v := h[len(h)-1]
	h = h[:len(h)-1]
	return v
}

在这种情况下，Push 和 Pop 方法对 h 的修改都不会反映到原始的 hp 对象上。具体来说：

Push 方法内部的 append 操作只修改了 h 的副本，并没有修改原始的 hp 对象。
Pop 方法内部对 h 的切片操作同样只修改了副本。

Go中的container/heap包是面向对象的风格，主要通过传入一个实现了 heap.Interface 接口的类型实例，以确保传入的类型具有堆操作所需的方法。这样可以在没有泛型的情况下自定义堆的类型和行为。

继承与重写

heap.Interface 定义了一个堆需要实现的基本方法：

type Interface interface {
	sort.Interface
	Push(x any) // add x as element Len()
	Pop() any   // remove and return element Len() - 1.
}

heap.Interface继承了sort.Interface，这是标准库中用于排序的接口

type Interface interface {
	// Len is the number of elements in the collection.
	Len() int

	// Less reports whether the element with index i
	// must sort before the element with index j.
	//
	// If both Less(i, j) and Less(j, i) are false,
	// then the elements at index i and j are considered equal.
	// Sort may place equal elements in any order in the final result,
	// while Stable preserves the original input order of equal elements.
	//
	// Less must describe a transitive ordering:
	//  - if both Less(i, j) and Less(j, k) are true, then Less(i, k) must be true as well.
	//  - if both Less(i, j) and Less(j, k) are false, then Less(i, k) must be false as well.
	//
	// Note that floating-point comparison (the < operator on float32 or float64 values)
	// is not a transitive ordering when not-a-number (NaN) values are involved.
	// See Float64Slice.Less for a correct implementation for floating-point values.
	Less(i, j int) bool

	// Swap swaps the elements with indexes i and j.
	Swap(i, j int)
}

所以我们传入的实例至少需要实现Push, Pop, Len, Less, Swap这五个方法

对于Int类型，我们可以让我们自己写的堆继承sort.IntSlice

// IntSlice attaches the methods of Interface to []int, sorting in increasing order.
type IntSlice []int

func (x IntSlice) Len() int           { return len(x) }
func (x IntSlice) Less(i, j int) bool { return x[i] < x[j] }
func (x IntSlice) Swap(i, j int)      { x[i], x[j] = x[j], x[i] }

这样就只用写Push和Pop方法了

// MinHeap 继承自IntSlice, 而IntSlice默认是用Less方法，继承了IntSlice中的Less方法，所以默认创建的是小顶堆
type MinHeap struct {
	sort.IntSlice
}

func (h *MinHeap) Push(v any) {
	h.IntSlice = append(h.IntSlice, v.(int))
}

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

但要注意这样实现的堆是最小堆(Min_Heap)，因为Less方法中写的是x[i] < x[j]，那么Less如何理解呢？可以认为x[i]和x[j]是数组中前后相邻的两个数，而Less要实现的效果就是前面的x[i]小于后面的x[j]，这样就是升序，也就是最小堆。

这种情况下，想要实现大顶堆就需要额外重写Less方法

// Less 重写IntSlice中的Less方法
func (h *MaxHeap) Less(i, j int) bool {
	return h.IntSlice[i] > h.IntSlice[j]
}

这样GoLand旁边也会显示一个蓝色的O表示重写

上浮与下沉

这时你可能会有疑问，我都把Push和Pop写好了，不能直接调我写的Push和Pop吗？

不可以，必须要用heap.Push和heap.Pop, 虽然heap.Push和heap.Pop底层是调用我们写的Push和Pop，
接调用我们写的Push和Pop，并不是按照堆的方式添加和弹出，仅仅只是简单的添加和弹出

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

我们相当于只用告诉堆添加和弹出的方式就行了，而heap包下主要做的是上浮和下沉的操作(up and down)

可以看到这里的down用到了我们写的Less方法，而这里的语义就不仅仅是数组中前面的元素应该比后面大还是小，而且前面的优先级比后面大

这段源码中，只有j2比j1位置的元素的优先级高，我们才让j = j2，只有j比i的优先级高，也就是比原先的parent高，才往下作交换并循环

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
}

这里我们写的代码为什么是获取IntSlice的最后一个元素并移除，而不是第一个呢

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

我们来看看Pop的底层实现

func Pop(h Interface) any {
	n := h.Len() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()
}

主要是将0号元素和最后一个元素对调了，这时只需要将最后一个元素弹出就可以获取优先级最高的元素了

本文采用署名-非商业性使用-相同方式共享 4.0 国际许可协议，转载请注明出处。