Go语言高级特性之优先队列|青训营

Go语言高级特性之优先队列

本篇博文的起因是：今天在刷力扣的每日一题时，在题解中看到了采用优先队列的解法，于是一不做二不休，顺便把自己的学习过程给记录下来。

Heap是一种数据结构，常用于实现优先队列，本文着重讨论Go中的堆结构和优先队列的实现。

什么是Heap

Heap（堆）是一种数据结构，其中包含一个特殊的根节点，且每个节点的值都不小于（或不大于）其所有的子节点的值。该种结构常用于实现优先级队列。

Heap的数据结构

Heap可以通过一个数组来实现，该数组应满足一下条件：

• 和二叉搜索树不同，不需要满足左子树的值小于父节点的值，右子树的值大于父节点的值

堆的一些列节点会按照某种特定的顺序进行排列。这样的顺序可以是最小的元素在最前面，也可以是最大的元素在最前面。这一顺序必须

满足父节点小于（或大于）其所有的子节点

• 堆中的元素并不一定是满的，即堆不一定是满二叉树

堆具有以下属性：

• 任何节点都小于（或大于）其所有后代，并且最小元素（或最大元素）位于堆的根（堆有序性）。

• 堆始终是一棵完整的树。即各级节点都填充除底层以外的元素，并且底层尽可能从左到右填充。

根节点最大的堆称为最大堆或大根堆，根节点最小的堆称为最小堆或小根堆。

由于堆是完全二叉树，因此它们可以表示为顺序数组，如下所示。

如何实现优先队列

优先队列是一种数据结构，其中每个元素都有一个优先级，优先级高的元素在前面，优先级相同时按照插入顺序排列。可以使用堆来实现

优先队列。实现优先队列的关键是将一个元素添加到队列中，并保持队列中的元素有序。如果使用数组来存储元素，需要频繁对数组进行

调整，时间复杂度是O(n)，不够高效。如果使用堆来存储元素，则可以在插入时进行堆化，时间复杂度是O(nlogn)。

在堆中，节点的位置与它们在数组中的位置有一定的关系。例如，根节点位于数组的第一个元素，其他节点依次排列。左子节点位于(2i),

右子节点位于(2i+1)，父节点位于(i/2)。这个关系可以方便地实现在数组上进行堆化的操作。

具体内容可以查看堆排序相关文章

优点和缺点

优点

• 简单高效：优先级队列的实现较为简单，查找和插入等操作都可以在 O(log(n))的时间复杂度内完成，所以在实现简单的情况下，可以极大提高程序性能。

• 优先级：优先级队列可以根据任务或者事件的优先级，对其按照优先级大小进行排序，并在需要的时候依次处理。

缺点

• 空间占用：优先级队列需要占用额外的内存空间，以存储任务和事件的优先级信息。

• 任务时效性：当优先级较高的任务过多时，可能会导致低优先级任务的响应延迟，从而影响任务的时效性。

heap 标准库

heap.Interface

其实Go的标准库已经为我们定义好了heap相关的接口，让我们看一下接口内的方法：

type Interface interface {
	sort.Interface
	Push(x any) // 将x添加至元素Len()的位置
	Pop() any   // 移除并且返回元素Len()-1
}

实现它的任何类型都可以用作具有以下不变量的最小堆（在调用Init()后或数据为空或排序时建立）。

注意，这个接口中的Push()和Pop()是由包堆的实现调用的。要从堆中添加和删除内容，请使用heap.Push()和heap.Pop()。

sort.Interface

heap.Interface中还包含了sort.Interface接口

type Interface interface {
        Len() int   // Len是集合中的元素个数。
        Less(i, j int) bool
        Swap(i, j int) // Swap交换索引i和索引j的元素
}

该接口意如其名，是为了进行排序而提供的接口。sort包下的方法通过整数索引来引用基础集合的元素。

Len()和Swap()函数的含有已在注释中写明，在这里重点讨论一下Less()函数。

Less()函数用来判断索引i的元素是否小于索引j的。

如果Less(i, j)和Less(j, i)均返回false，那么认为这两个索引对应的元素相等。

Sort在最终结果中按任意顺序放置相等的元素，而Stable(稳定排序)保留相等元素的原始输入顺序。

Less()函数必须具有传递性：

• 若Less(i, j)和Less(j, k)均为true，则Less(i, k)也为true
• 若Less(i, j)和Less(j, k)均为false，则Less(i, k)也为false

注意，当涉及非数字，或者浮点型（float32或float64）比较不具有传递性。

初始化堆

Init函数建立此包中其他程序所需的堆不变量（heap invariants）。

Init()对于堆不变量是幂等的(可被多次调用)，并且可以在堆不变量失效时调用。

该函数的时间复杂度为O(n),n为堆中元素的个数。

func Init(h Interface) {
	// heapify
	n := h.Len()
	for i := n/2 - 1; i >= 0; i-- {
		down(h, i, n)
	}
}

该步骤就是堆排序中对数组进行堆化的过程，即对每个元素进行一次**下沉（down）**操作。

操作堆

• Push()：插入元素。
• Pop()：移除并返回堆顶元素。
• Remove()：移除并返回第i个元素。
• Fix()：修复堆。

Push()

Push()将元素x插入堆中。

该方法的时间复杂度为O(log n),n为堆中元素个数。

func Push(h Interface, x any) {
	h.Push(x)
	up(h, h.Len()-1)
}

从源码中可以看出，该函数先调用了heap.Interface接口的Push()函数将x元素放置到堆的最后，然后对该元素进行**上浮（up）**操作。

Pop()

Pop()从堆中移除并返回最小元素（根据Less的规则）。也就是堆顶元素。

Pop()的时间复杂度是O(log n)，n为堆中的元素个数。

Pop()等价于Remove(h, 0)。

func Pop(h Interface) any {
	n := h.Len() - 1
	h.Swap(0, n)
	down(h, 0, n)
	return h.Pop()
}

从源码可以看出，实际的Pop步骤是将第一个元素与最后一个元素进行互换，然后对一个元素进行下沉，最后移除并返回最后一个元素。

Remove()

Remove()从堆中移除并返回第i个元素。

Remove()的时间复杂度是O(log n)，n为堆中的元素个数。

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()

Fix()在第i个元素的值改变后重新建立堆排序。

更改第i个元素的值，然后调用Fix()相当于调用Remove(h, i)，然后再Push新值，但成本比调用Remove(h, i)低。

Fix()的时间复杂度为O(log n)，n为堆中的元素个数。

func Fix(h Interface, i int) {
	if !down(h, i, h.Len()) {
		up(h, i)
	}
}

up and down

up和down都是未导出函数，即不能被外部调用的函数。

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
}

简单实现

这里以int类型为例，我们先定义堆的底层数组

type IntHeap []int

然后我们需要实现heap.Interface中的5个函数

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{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

需要注意的是这里的Push和Pop函数都需要传入指针参数，因为这两个函数都会实际修改切片中的元素。

然后我们写一个main函数来测试一下：

func main() {
	h := &IntHeap{6, 2, 4, 1}
	heap.Init(h)
	heap.Push(h, 3)
	fmt.Printf("堆顶元素: %d\n", (*h)[0])
	for h.Len() > 0 {
		fmt.Printf("%d ", heap.Pop(h))
	}

	// 输出:
	//堆顶元素: 1
	//1 2 3 4 6

}

优先队列

官方也在该包下面为我们实现了一个PriorityQueue的测试用例。

测试用的结构体为：

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

然后就是上面相同的套路，定义堆的底层结构体数组，并且实现heap.Interface中的方法：

type PriorityQueue []*Item

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

func (pq PriorityQueue) Less(i, j int) bool {
	// 此处我们返回的是优先级更高的Item
	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 any) {
	n := len(*pq)
	item := x.(*Item)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() any {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil  // 防止内存泄漏
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

官方在这里为我们多提供了一个函数，用于更新堆内某一Item的属性值：

// 修改队列中Item的priority和value
func (pq *PriorityQueue) update(item *Item, value string, priority int) {
	item.value = value
	item.priority = priority
	heap.Fix(pq, item.index)
}

测试用例：

func Example_priorityQueue() {
	// Some items and their priorities.
	items := map[string]int{
		"banana": 3, "apple": 2, "pear": 4,
	}

	// 创建优先队列，将元素放入，并且建立优先队列不变量
	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)

	// 将元素插入，并修改其属性
	item := &Item{
		value:    "orange",
		priority: 1,
	}
	heap.Push(&pq, item)
	pq.update(item, item.value, 5)

	// 将所有元素pop出队列，按照优先级递减顺序
	for pq.Len() > 0 {
		item := heap.Pop(&pq).(*Item)
		fmt.Printf("%.2d:%s ", item.priority, item.value)
	}
    // 输出：
	// 05:orange 04:pear 03:banana 02:apple
}

优先队列Plus

上面我们展示了优先队列的一个简单实现，下面我们增加亿点点难度，放到实际的生产环境中，即存在生产者和消费者的场景中来讨论。

生产者

对于生产者来说，他只需要推送一个任务及其优先级过来，咱们就得根据优先级处理他的任务。

由于，我们不大好判断，到底会有多少种不同的优先级传过来，也无法确定，每种优先级下有多少个任务要处理，所以，我们可以直接使

用heap存储task

消费者

对于消费者来说，他需要获取优先级最高的任务进行消费。使用heap中的Pop方法 取出优先级最高的任务即可

数据结构

1）Item（PriorityQueue中实际存储的对象，对task进行包装）

type queue []*Item//Item对象的底层数组

type Item struct {
	Key      string
	Value    interface{}
	Priority int64
	index    int
}

2）PriorityQueue

type PriorityQueue struct {
	data    queue
	dataMap map[string]*Item
	lock    sync.RWMutex
}//队列内部采用切片和以Key为索引的map来维护元素

func NewPriorityQueue() *PriorityQueue {//创建一个PriorityQueue并初始化切片和map
	pq := PriorityQueue{
		data:    make(queue, 0),
		dataMap: make(map[string]*Item),
	}
	heap.Init(&pq.data)
	return &pq
}

3）优先队列对象

type PriorityQueueTask struct {//对PriorityQueue的封装，同时包含监听的channel
    mLock    sync.Mutex // 互斥锁，queues和priorities并发操作时使用，当然针对当前读多写少的场景，也可以使用读写锁
    pushChan chan *task // 推送任务管道
    pq       *PriorityQueue
}

4）任务对象

type task struct {
    priority int64 // 任务的优先级
    value    interface{}
    key      string
}

实现接口

为了实现堆结构，需要实现heap.Interface接口

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]
	q[i].index = i
	q[j].index = j
}

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

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

func (pq *PriorityQueue) Len() int {
	pq.lock.RLock()
	defer pq.lock.RUnlock()
	return pq.data.Len()
}

func (pq *PriorityQueue) Pop() (*Item, error) {
	pq.lock.Lock()
	defer pq.lock.Unlock()

	if pq.data.Len() == 0 {
		return nil, ErrEmpty//返回queue为空的err
	}

	item := heap.Pop(&pq.data).(*Item)
	delete(pq.dataMap, item.Key)
	return item, nil
}

func (pq *PriorityQueue) Push(i *Item) error {
	if i == nil || i.Key == "" {
		return errors.New("error adding item: Item Key is required")
	}

	pq.lock.Lock()
	defer pq.lock.Unlock()

	if _, ok := pq.dataMap[i.Key]; ok {
		return ErrDuplicateItem//返回Key重复的err
	}

    // 通过克隆防止修改直接影响到实际的Item
    //此处克隆结构体需import("github.com/mitchellh/copystructure")
	clone, err := copystructure.Copy(i)
	if err != nil {
		return err
	}

	pq.dataMap[i.Key] = clone.(*Item)
	heap.Push(&pq.data, clone)
	return nil
}

//额外提供一个根据Key来Pop的函数
func (pq *PriorityQueue) PopByKey(key string) (*Item, error) {
	pq.lock.Lock()
	defer pq.lock.Unlock()

	item, ok := pq.dataMap[key]
	if !ok {
		return nil, nil
	}

	itemRaw := heap.Remove(&pq.data, item.index)
	delete(pq.dataMap, key)

	if itemRaw != nil {
		if i, ok := itemRaw.(*Item); ok {
			return i, nil
		}
	}

	return nil, nil
}

初始化优先队列对象

在初始化对象时，我们通过NewPriorityQueueTask函数创建一个PriorityQueue对象赋值给pq，然后开启推送任务的管道。

func NewPriorityQueueTask() *PriorityQueueTask {
    pq := &PriorityQueueTask{
        pushChan: make(chan *task, 100),
        pq:       NewPriorityQueue(),
    }
    // 监听pushChan
    go pq.listenPushChan()
    return pq
}
func (pq *PriorityQueueTask) listenPushChan() {
    for {
        select {//如果pushChan被推送，泽则吧task封装成Item，Push进pq中
        case taskEle := <-pq.pushChan:
            pq.mLock.Lock()
            pq.pq.Push(&Item{Key: taskEle.key, Priority: taskEle.priority, Value: taskEle.value})
            pq.mLock.Unlock()
        }
    }
}

生产者推送任务

生产者向推送任务管道中推送新任务时，实际上是将一个 task 结构体实例发送到了管道中。在 task 结构体中，priority 属性表示这个任

务的优先级，value 属性表示这个任务的值，key 属性表示这个任务的键。

func (pq *PriorityQueueTask) Push(priority int64, value interface{}, key string) {
    pq.pushChan <- &task{
        value:    value,
        priority: priority,
        key:      key,
    }
}

消费者消费队列

消费者从队列中取出一个任务，然后进行相应的操作。在这段代码中，消费者轮询获取最高优先级的任务。如果没有获取到任务，则继续

轮询；如果获取到了任务，则执行对应的操作。在这里，执行操作的具体形式是打印任务的Value、优先级等信息。

func (pq *PriorityQueueTask) Consume() {
    for {
        task := pq.Pop()
        if task == nil {
            // 未获取到任务，则继续轮询
            time.Sleep(time.Millisecond)
            continue
        }
        // 获取到了任务，就执行任务
        fmt.Println("推送任务的编号为：", task.Value)
        fmt.Println("推送的任务优先级为：", task.Priority)
        fmt.Println("============")
    }
}

测试用例

接下来我们写一个测试用例:

func TestQueue(t *testing.T) {
	defer func() {
		if err := recover(); err != nil {
			fmt.Println(err)
		}
	}()
	pq := NewPriorityQueueTask()
	// 在这里随机生成一些优先级任务
	for i := 0; i < 100; i++ {
		a := rand.Intn(1000)
		go func(a int64) {
			pq.Push(a, a, strconv.Itoa(int(a)))
		}(int64(a))
	}
	// 这里会阻塞，消费者会轮询查询任务队列
	pq.Consume()
}

结语

Go语言中heap的实现采用了一种 “模板设计模式”，用户实现自定义堆时，只需要实现heap.Interface接口中的函数，然后应用heap.Push、heap.Pop等方法就能够实现想要的功能。本文中的部分代码实现来自于vault/sdk/queue/priority_queue.go at main · hashicorp/vault (github.com)