Go_Src 详解-Heap

发布日期: 2024-05-24

更新日期: 2024-05-27

文章字数: 1.2k

阅读时长: 5 分

源码中文注释版


// heap 包为任何实现了 heap.Interface 的类型提供堆操作。堆是一棵树，其中每个节点都是其子树中的最小值节点。
//
// 树中的最小元素是根，位于索引 0。
//
// 堆是实现优先队列的常见方式。要构建优先队列，请使用（负）优先级作为 Less 方法的排序方式实现 Heap 接口，
// 因此 Push 添加项目，而 Pop 从队列中删除最高优先级的项目。示例包括这样的实现；文件 example_pq_test.go 包含完整的源代码。
package heap

import "sort"
 
// Interface 类型描述了使用此包中的例程的类型的要求。
// 任何实现它的类型都可以用作具有以下不变量的最小堆（在调用 Init 后或者数据为空或已排序时建立）：
//
//  !h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
//
// 请注意，此接口中的 Push 和 Pop 是供包 heap 的实现调用的。要从堆中添加和删除内容，请使用 heap.Push 和 heap.Pop。
type Interface interface {
    sort.Interface
    Push(x any) // add x as element Len()
    Pop() any   // remove and return element Len() - 1.
}

// Init 建立此包中其他例程所需的堆不变量。
// Init 对于堆不变量是幂等的，并且可以在堆不变量可能失效时调用。
// 复杂度为 O(n)，其中 n = h.Len()。
func Init(h Interface) {
    // heapify
    n := h.Len()
    for i := n/2 - 1; i >= 0; i-- {
       down(h, i, n)
    }
}

// Push 将元素 x 推入堆中。
// 复杂度为 O(log n)，其中 n = h.Len()。
func Push(h Interface, x any) {
    h.Push(x)
    up(h, h.Len()-1)
}

// Pop 从堆中删除并返回最小元素（根据 Less）。
// 复杂度为 O(log n)，其中 n = h.Len()。
// Pop 等效于 Remove(h, 0)。
func Pop(h Interface) any {
    n := h.Len() - 1
    h.Swap(0, n)
    down(h, 0, n)
    return h.Pop()
}

// Remove 从堆中删除并返回索引 i 处的元素。
// 复杂度为 O(log n)，其中 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 在索引 i 处的元素更改其值后重新建立堆排序。
// 更改索引 i 处元素的值，然后调用 Fix 等效于调用 Remove(h, i) 然后 Push 新值，但开销更小。
// 复杂度为 O(log n)，其中 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
}

使用要点：

手动实现 heap.Interface 接口
具体使用的时候，是使用 heap 的函数，而不是自己的定义的接口的函数
- right：heap.Push(&q, node.Next)
- wrong: q.Push(node.Next)
编写 POP 函数的时候，要注意释放对应的空间，避免内存泄露（参考 Go-实战之 Go Slice 的坑）

Demo1

小根堆

package main

import "container/heap"


// primaryQueue 是 ListNode 指针的最小堆/优先队列
type primaryQueue []*ListNode

// 约束：primary_queue 实现 heap.Interface 接口
var _ heap.Interface = (*primaryQueue)(nil)

// Len 方法返回堆中的元素个数
func (p *primaryQueue) Len() int {
	return len(*p)
}

// Less 方法定义了堆的排序规则
func (p *primaryQueue) Less(i, j int) bool {
	return (*p)[i].Val < (*p)[j].Val
}

// Swap 方法交换了堆中的两个元素
func (p *primaryQueue) Swap(i, j int) {
	(*p)[i], (*p)[j] = (*p)[j], (*p)[i]
}

// Push 方法向堆中添加元素
func (p *primaryQueue) Push(x any) {
	*p = append(*p, x.(*ListNode))
}

// Pop 方法从堆中弹出元素
func (p *primaryQueue) Pop() any {
	v := (*p)[len(*p)-1]
    (*p)[len(*p)-1]=nil // 释放对应的内存,避免内存泄漏
	*p = (*p)[:len(*p)-1]
	return v
}

func mergeKLists(lists []*ListNode) *ListNode {
	q := make(primaryQueue, 0, len(lists))
	for _, head := range lists {
		if head != nil {
			q = append(q, head)
		}
	}
	heap.Init(&q) // 堆化

	dummy := &ListNode{} // 哨兵节点，作为合并后链表头节点的前一个节点
	cur := dummy
	for len(q) > 0 { // 循环直到堆为空
		node := heap.Pop(&q).(*ListNode) // 剩余节点中的最小节点
		if node.Next != nil {            // 下一个节点不为空
			heap.Push(&q, node.Next) // 将下一个节点推入堆中
		}
		cur.Next = node // 合并到新链表中
		cur = cur.Next  // 准备合并下一个节点
	}
	return dummy.Next // 哨兵节点的下一个节点就是新链表的头节点
}