用go实现一个优先队列

上篇文章中提到了一个优先队列，正好拿来说一下。

优先队列有别于普通的队列，普通的队列是先进先出，而优先队列是优先级高的先出，优先级低的后出。

优先队列与普通队列的区别在于“谁出”，优先队列可以使用双链表来实现，只需要找到优先级最高的节点然后删除返回即可，时间复杂度为 \(O(n)\)。那还能不能优化呢？当然可以，可以使用二叉堆这种数据结构，其构建时间复杂度为 \(O(nlogn)\)。

具体来说，我们需要最大二叉堆这个结构，这样才能高效地删除最大优先级的元素。

先定义一个元素：

type element struct {
	value    interface{}
	priority int
}

定义队列，要实现大小比较方法和 heap 的接口：

type queue []element

// Len、Less、Swap 是排序方法的定义

// 队列长度
func (q queue) Len() int {
	return len(q)
}
// 我们设定最大堆
func (q queue) Less(i, j int) bool {
	return q[i].priority > q[j].priority
}
func (q queue) Swap(i, j int) {
	q[i], q[j] = q[j], q[i]
}

// 实现二叉堆的接口

// 添加元素
func (q *queue) Push(v interface{}) {
	*q = append(*q, v.(element))
}

// 弹出元素
func (q *queue) Pop() interface{} {
	l := len(*q)
	ele := (*q)[l-1]
	*q = (*q)[:l-1]
	return ele
}

我们对上面的 queue 封装下，仅暴露出必要的接口：

// 对 queue 封装一层，只保留必要的对外api
type PriorityQueue struct {
	q *queue
}

// 初始化
func (pq *PriorityQueue) Init() {
	tmpQ := make(queue, 0)
	pq.q = &tmpQ
}

// 添加数据
func (pq *PriorityQueue) Push(v interface{}, priority int) {
	heap.Push(pq.q, element{ // 使用 heap.Push 二叉堆会排序
		value:    v,
		priority: priority,
	})
}

// 弹出最大优先级数据
func (pq *PriorityQueue) Pop() interface{} {
	return heap.Pop(pq.q).(element).value
}

// 长度
func (pq *PriorityQueue) Len() int {
	return pq.q.Len()
}

// 是否为空
func (pq *PriorityQueue) IsEmpty() bool {
	return pq.q.Len() == 0
}

测试：

func main() {
	var pq PriorityQueue
	pq.Init()
	pq.Push("a", 4)
	pq.Push("b", 5)
	pq.Push("c", 2)
	pq.Push("d", 9)
	pq.Push("e", 1)
	fmt.Printf("length:%v \n", pq.Len())

	fmt.Printf("pop element:%v, and result length:%v \n", pq.Pop(), pq.Len())
	fmt.Printf("pop element:%v, and result length:%v \n", pq.Pop(), pq.Len())
}

结果：

length:5
pop element:d, and result length:4
pop element:b, and result length:3