go语言container/heap 源码解读与应用阅读笔记

golang container/heap源码阅读笔记

1. 源码解析

type Interface interface {
	sort.Interface
    Push(x interface{}) // add x as element Len()  将x作为第len()个元素加入堆中
	Pop() interface{}   // remove and return element Len() - 1. 
}

首先定义heap的接口，其中 sort.Interface包括一下三个接口，任何类型，只要实现了这五个接口，就是一个heap

Interface interface {
    // Len is the number of elements in the collection. 集合内的元素个数
    Len() int
    // Less reports whether the element with
    // index i should sort before the element with index j.
    // 返会索引为i的元素是否应该在j的前面
    Less(i, j int) bool
    // Swap swaps the elements with indexes i and j.
    Swap(i, j int)
}

堆的逻辑结构是满二叉树，堆的常用的存储结构有顺序存储、链表存储，go的官方源码采用的是顺序结构。

以小根堆为例说明

down()函数是从上到下调整的过程，就是从节点i0开始，将其与子节点中较小的节点交换，直到i最后的位置为叶节点或者以i为父节点的子树满足小根堆的要求。

如果发生了交换，i就比i0大，否则 i等于i0 返回值 i>i0表示是否发生的交换

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
		}
        // 如果j1大于n,说明索引i没有孩子节点，说明它本身就是孩子节点，此时已经
		j := j1 // left child
		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
            // 判断右孩子节点是否存在，且 如果小于j1，就将j2赋值给j
			j = j2 // = 2*i + 2  // right child
		}
        // 此时得到的j为i的两个孩子节点较小的索引
		if !h.Less(j, i) {
            // 判断 孩子节点与父节点是否需要交换，如果不需要，说明已经满足堆的要求（小根或者大根）
			break
		}
		h.Swap(i, j)
        // 交换父节点和子节点
		i = j
        
	}
	return i > i0
}

up()函数是从下向上调整的过程，将j节点与其父节点比较，若满足小根堆的根值小于孩子值，就交换。

注意，up()函数不考虑兄弟节点，只要自己比父亲小，就可以交换。

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

结合上图 init()函数就是从非叶节点开始，逆序调整节点，构建堆heap

// The complexity is O(n) where 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入堆，入堆操作有两步，先讲元素x放到树的最末尾，然后向上调整。

// Push pushes the element x onto the heap.
// The complexity is O(log n) where n = h.Len().
func Push(h Interface, x interface{}) {
	h.Push(x)
	up(h, h.Len()-1)
}

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) interface{} {
	n := h.Len() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()
}

Remove()函数是将索引为i的元素从堆中移除，以pop()函数的处理手法相同，将索引为i的元素与堆最末尾的元素交换，然后判断新换上来的元素是否需要调整(down()函数返回bool值得同时，也调整了)，如果没有向下调整，说明被换上来的元素是以以这个位置为根的最小值，需要向上尝试调整

// 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) interface{} {
	n := h.Len() - 1
	if n != i {
		h.Swap(i, n)
		if !down(h, i, n) {
			up(h, i)
		}
	}
	return h.Pop()
}

fix()函数是在改变堆中某个位置的值后，进行堆的再建，与remove()函数中 被移除元素与末尾元素交换后的处理流程一样。

如果该元素向下调整了，说明修改后的值比原来的值大，不需要向上调整。

如果该元素没有向下调整了，说明修改后的值比原来的值小，需要向上调整。

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

2. 堆的应用一（小根堆、大根堆）

package main

import (
	"container/heap"
	"fmt"
)

type Intheap []int

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

func (h *Intheap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *Intheap) Pop() interface{} {
	n := len(*h)
	x := (*h)[n-1]
	*h = (*h)[:n-1]
	return x
}

func main() {

	nums := &Intheap{2, 5, 3, 9, 4, 6, 8}
	heap.Init(nums)
	fmt.Println(nums)// [2 4 3 9 5 6 8]
	heap.Push(nums, 1)
	fmt.Println(nums) //[1 2 3 4 5 6 8 9]
	heap.Push(nums, 3)
	fmt.Println(nums)// [1 2 3 3 5 6 8 9 4]

	heap.Remove(nums, 4)
	fmt.Println(nums)//[1 2 3 3 4 6 8 9]

	(*nums)[1] = 10
	fmt.Println(nums)//[1 10 3 3 4 6 8 9]
	heap.Fix(nums, 1)
	fmt.Println(nums)//[1 3 3 9 4 6 8 10]

	x:=heap.Pop(nums)//1
	fmt.Println(x)

}

实现heap的五个接口即可。

上面实现的是小根堆，注意到less()函数，只需要将less函数中的小于号改为大于号即可实现大根堆

less()函数的含义是，如果返回值为真，就需要交换i和j的位置

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

3. 堆的应用二 优先队列

import (
	"container/heap"
	"fmt"
)


// An Item is something we manage in a priority queue.
type Item struct {
	value    string // The value of the item; arbitrary.
	priority int    // The priority of the item in the queue.
	// The index is needed by update and is maintained by the heap.Interface methods.
	index int // The index of the item in the heap.
}

// A PriorityQueue implements heap.Interface and holds Items.
type PriorityQueue []*Item

func (pq PriorityQueue) Len() int { return len(pq) }

func (pq PriorityQueue) Less(i, j int) bool {
	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
	return pq[i].priority > pq[j].priority
}

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

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

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil  // avoid memory leak
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

// update modifies the priority and value of an Item in the queue.
func (pq *PriorityQueue) update(item *Item, value string, priority int) {
	item.value = value
	item.priority = priority
	heap.Fix(pq, item.index)
}

// This example creates a PriorityQueue with some items, adds and manipulates an item,
// and then removes the items in priority order.
func main() {
	// Some items and their priorities.
	items := map[string]int{
		"banana": 3, "apple": 2, "pear": 4,
	}

	// Create a priority queue, put the items in it, and
	// establish the priority queue (heap) invariants.
	pq := make(PriorityQueue, len(items))
	i := 0
	for value, priority := range items {
		pq[i] = &Item{
			value:    value,
			priority: priority,
			index:    i,
		}
		i++
	}
	heap.Init(&pq)

	// Insert a new item and then modify its priority.
	item := &Item{
		value:    "orange",
		priority: 1,
	}
	heap.Push(&pq, item)
	pq.update(item, item.value, 5)

	// Take the items out; they arrive in decreasing priority order.
	for pq.Len() > 0 {
		item := heap.Pop(&pq).(*Item)
		fmt.Printf("%.2d:%s ", item.priority, item.value)
	}
	// Output:
	// 05:orange 04:pear 03:banana 02:apple
}

与大根堆小根堆不同的是每个元素是个item结构体，结构体内包含值、优先级、索引等

其中,

1. less函数，比较每个item的优先级来确认是否需要交换

2. swap函数，如何交换元素，还要改变索引