[golang] 4 - container/heap

Posted by Dongbo on March 9, 2022

golang里要使用优先队列略微麻烦,需要自己实现 heap.Interface 下的几个接口,比起 c++ 的容器使用起来更麻烦一些。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
type rank struct {
	score, index int
}

type priorityQueue []rank

func (q *priorityQueue) Push(x interface{}) {
	*q = append(*q, x.(rank)) // note: 需要赋值给指针否则无法更新实际的heap
}

func (q *priorityQueue) Pop() interface{} {
	pq := *q
	size := len(pq)
	x := pq[size - 1]
	*q = pq[:size - 1]  // 同上
	return x
}

func (q *priorityQueue) Less(i, j int) bool { // 定义比较函数
	pq := *q
	return pq[i].score > pq[j].score 
}

func (q *priorityQueue) Swap(i, j int) {
	pq := *q
	pq[i], pq[j] = pq[j], pq[i]
}

func (q *priorityQueue) Len() int {
	return len(*q)
}

默认是最小堆(如果Less()函数是小于的话),上面实现中为大于,表示实现的是按照 rank.score 排序的最大堆。另外要记得 c++ 中默认是最大堆。

调用示例:

1
2
3
4
5
6
7
8
9
10
11
func TestPQ(t *testing.T) {
	score := []int{1,5,12,6,4,8,36,20}
	pq := priorityQueue{}
	for i := range score {
		pq = append(pq, rank{score: score[i], index: i})
	}
	heap.Init(&pq)
	for pq.Len() > 0 {
		t.Log(heap.Pop(&pq))
	}
}

The End