golang之heap学习

测试环境

go version go1.11.2 windows/amd64

概述

heap, 即堆, 是一种用数组实现的完全二叉树.
堆有大根堆和小根堆, 分别是说: 对应的二叉树的树根结点的键值是所有堆节点键值中最大/小者。
为了方便叙述, 以小根堆为例说下概念.
堆常见的实现方法: 在数组(arr)中存储若干个元素, 如果

• 数组第一个元素为哨兵(如果是小根堆, 则arr[0]可以放一个很小的值, 比其他元素都小), 那么任意一个实际的节点的索引j和其父节点索引i有如下关系(L, R分别表示左右):
  $j_L=2i$, $j_R=2i+1$;

• 数组中所有位置都存放实际数据, 那么: $j_L=2i+1,j_R=2i+2$;

golang中的heap用的是上述第二种存储方式, 因为数据类型是interface{}, 所以没办法放哨兵, 当然必要性也不大, 除了计算上的方便.
heap源码位置: $GOROOT/src/container/heap/heap.go

例子演示

下面通过一些例子来演示heap的用法
go源码带的IntHeap($GOROOT/src/container/heap/example_intheap_test.go)已经演示了最基本的用法. 我们自己写个MonthHeap试试.

MonthHeap

MonthHeap是这样的, 它里边存放月份的英文简写, 序号较前的月份应该在堆顶.
完整源代码在这里;

package month_heap

import (
	"strings"
)

var monthMap map[string]int = make(map[string]int)

func init() {
	months := strings.Split("Jan, Feb, Mar, Apr, May, June, July, Aug, Sep, Oct, Nov, Dec", ",")
	for i, v := range months { // 将月份简写与序号对应
		monthMap[strings.TrimSpace(v)] = i + 1
	}
}

type MonthHeap []string

// 实现heap.Interface需要的5个方法
func (m MonthHeap) Len() int {
	return len(m)
}

// 注意这里: 比较的是哪个月份在前
func (m MonthHeap) Less(i, j int) bool {
	return monthMap[m[i]] < monthMap[m[j]]
}

func (m MonthHeap) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

// Push和Pop都要用指针接收者, 因为要在函数内修改slice
func (m *MonthHeap) Push(x interface{}) {
	*m = append(*m, x.(string))
}

func (m *MonthHeap) Pop() interface{} {
	old := *m
	n := len(old)
	x := old[n-1]
	*m = old[0 : n-1]
	return x
}

测试如下:

func MonthHeapTest() {
	h := &month_heap.MonthHeap{"Jan", "Feb", "Mar"}
	heap.Init(h)
	heap.Push(h, "May")
	heap.Push(h, "Apr")
	// first month: Jan
	fmt.Println("first month:", (*h)[0])

	// 输出: Jan	Feb	Mar	Apr	May
	for h.Len() > 0 {
		fmt.Printf("%s\t", heap.Pop(h)) // 注意不是h.Pop()
	}
	fmt.Println()
}

优先级队列priority_queue

完整代码在这里

package priority_queue

import "container/heap"

type Item struct {
	Value    string // 值
	Priority int    // 优先级
	Index    int    // 元素在堆中的索引
}

type PriorityQueue []*Item

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

// 注意这里是优先级大的在前
func (pq PriorityQueue) Less(i, j int) bool {
	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{}) {
	item := x.(*Item)
	n := len(*pq)
	item.Index = n
	*pq = append(*pq, item)
}

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

	// 为了安全起见? 应该是防止直接用这个值插入或者别的,
	// 将index设置为一个越界值, 可以防止误操作
	x.Index = -1
	*pq = old[0 : n-1]
	return x
}

// 更新queue中一个Item的Priority和value, 将将Item调整到queue中适当的位置(满足堆的特征)
// 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

	// Fix操作比 先Remove再Push要快
	heap.Fix(pq, item.Index)
}

测试代码:

func PriorityQueueTest() {
	items := map[string]int{
		"AA": 5,
		"BB": 8,
		"CC": 3,
	}
	pq := make(priority_queue.PriorityQueue, len(items))
	i := 0

	for value, priority := range items {
		pq[i] = &priority_queue.Item{
			Value:    value,
			Priority: priority,
			Index:    i,
		}
		i++
	}
	heap.Init(&pq)

	item := priority_queue.Item{
		Value:    "DD",
		Priority: 1,
	}
	heap.Push(&pq, &item)
	pq.Update(&item, "EE", 99) // 修改名字并调整优先级

	// EE:99	BB:8	AA:5	CC:3
	for pq.Len() > 0 {
		x := heap.Pop(&pq).(*priority_queue.Item)
		fmt.Printf("%s:%d\t", x.Value, x.Priority)
	}
	fmt.Println()
}

欢迎补充指正!