实现优先级队列和并发安全优先级队列 (#110)

基于小顶堆和切片的实现

.CHANGELOG.md

@@ -1,6 +1,7 @@
 # 开发中
 - [atomicx: 泛型封装 atomic.Value](https://github.com/gotomicro/ekit/pull/101)
 - [queue: API 定义](https://github.com/gotomicro/ekit/pull/109)
+- [queue: 基于堆和切片的优先级队列](https://github.com/gotomicro/ekit/pull/110)
 
 # v0.0.4
 - [slice: 重构 index 和 contains 的方法，直接调用对应Func 版本](https://github.com/gotomicro/ekit/pull/87)

internal/queue/priority_queue.go

@@ -0,0 +1,136 @@
+// Copyright 2021 gotomicro
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package queue
+
+import (
+	"errors"
+
+	"github.com/gotomicro/ekit/internal/slice"
+
+	"github.com/gotomicro/ekit"
+)
+
+var (
+	ErrOutOfCapacity = errors.New("ekit: 超出最大容量限制")
+	ErrEmptyQueue    = errors.New("ekit: 队列为空")
+)
+
+// PriorityQueue 是一个基于小顶堆的优先队列
+// 当capacity <= 0时，为无界队列，切片容量会动态扩缩容
+// 当capacity > 0 时，为有界队列，初始化后就固定容量，不会扩缩容
+type PriorityQueue[T any] struct {
+	// 用于比较前一个元素是否小于后一个元素
+	compare ekit.Comparator[T]
+	// 队列容量
+	capacity int
+	// 队列中的元素，为便于计算父子节点的index，0位置留空，根节点从1开始
+	data []T
+}
+
+func (p *PriorityQueue[T]) Len() int {
+	return len(p.data) - 1
+}
+
+// Cap 无界队列返回0，有界队列返回创建队列时设置的值
+func (p *PriorityQueue[T]) Cap() int {
+	return p.capacity
+}
+
+func (p *PriorityQueue[T]) IsBoundless() bool {
+	return p.capacity <= 0
+}
+
+func (p *PriorityQueue[T]) isFull() bool {
+	return p.capacity > 0 && len(p.data)-1 == p.capacity
+}
+
+func (p *PriorityQueue[T]) isEmpty() bool {
+	return len(p.data) < 2
+}
+
+func (p *PriorityQueue[T]) Peek() (T, error) {
+	if p.isEmpty() {
+		var t T
+		return t, ErrEmptyQueue
+	}
+	return p.data[1], nil
+}
+
+func (p *PriorityQueue[T]) Enqueue(t T) error {
+	if p.isFull() {
+		return ErrOutOfCapacity
+	}
+
+	p.data = append(p.data, t)
+	node, parent := len(p.data)-1, (len(p.data)-1)/2
+	for parent > 0 && p.compare(p.data[node], p.data[parent]) < 0 {
+		p.data[parent], p.data[node] = p.data[node], p.data[parent]
+		node = parent
+		parent = parent / 2
+	}
+
+	return nil
+}
+
+func (p *PriorityQueue[T]) Dequeue() (T, error) {
+	if p.isEmpty() {
+		var t T
+		return t, ErrEmptyQueue
+	}
+
+	pop := p.data[1]
+	p.data[1] = p.data[len(p.data)-1]
+	p.data = p.data[:len(p.data)-1]
+	p.shrinkIfNecessary()
+	p.heapify(p.data, len(p.data)-1, 1)
+	return pop, nil
+}
+
+func (p *PriorityQueue[T]) shrinkIfNecessary() {
+	if p.IsBoundless() {
+		p.data = slice.Shrink[T](p.data)
+	}
+}
+
+func (p *PriorityQueue[T]) heapify(data []T, n, i int) {
+	minPos := i
+	for {
+		if left := i * 2; left <= n && p.compare(data[left], data[minPos]) < 0 {
+			minPos = left
+		}
+		if right := i*2 + 1; right <= n && p.compare(data[right], data[minPos]) < 0 {
+			minPos = right
+		}
+		if minPos == i {
+			break
+		}
+		data[i], data[minPos] = data[minPos], data[i]
+		i = minPos
+	}
+}
+
+// NewPriorityQueue 创建优先队列 capacity <= 0 时，为无界队列，否则有有界队列
+func NewPriorityQueue[T any](capacity int, compare ekit.Comparator[T]) *PriorityQueue[T] {
+	sliceCap := capacity + 1
+	if capacity < 1 {
+		capacity = 0
+		sliceCap = 64
+	}
+	return &PriorityQueue[T]{
+		capacity: capacity,
+		data:     make([]T, 1, sliceCap),
+		compare:  compare,
+	}
+}