实现一个后入先出的优先级队列
最近需要做个优先级队列来做任务,虽然 java.util中提供了 PriorityQueue ,但那个的应用场景和期望效果不太匹配。目标是做一个元素具有指定优先级属性,并根据不同的优先级数据按插入顺序的倒序执行的效果,并且支持取消未执行的任务。
执行顺序的原则:每个元素都有一个优先级(高、中、低),任务的执行顺序是:高优先级的执行顺序一定高于低优先级,同一优先级按插入顺序,后入先出。
涉及思路
封装一个新的数据结构最核心的一个步骤是选择底层的数据结构。比如 HashMap 底层是链表和红黑树、ArrayList 底层数据结构是数组等等。
优先级队列支持通过索引快速取消任务,那么这个数据结构就应该使用 Map 来保存键值对关系,使用 HashMap 来保存真实数据,HashMap 查询的时间复杂的为 log(n) ,也很适合快速取出(提高效率)。
根据优先级队列排序的特点,需要另一种数据结构来处理排序。有两种数据结构比较合适,数组和链表。数组支持快速索引,但需要设计扩容机制;链表则无需关注容量问题,虽然链表不支持快速索引,但任务执行每次都是取头节点,链表的头节点取出时间复杂度也是 O(1),仅在删除时需要进行遍历。
事实上,底层通过三个保存不同优先级的 ArrayList ,是一个很方便的选择,因为 ArrayList 具有自动扩容机制,只需按不同优先级的 ArrayList 取出任务进行消费即可。但这样的话未免过于简单。所以我选择用链表来实现。
通过以 Key-Value 的形式保存数据,PriorityQueue 类的定义大概是这样的:
class PriorityQueue<K, V>: Iterable<V> {
fun put(key: K?, value: V?, priority: Priority): Boolean
fun pop(): V?
fun peek(): V?
fun remove(key: K): Boolean
// ...
}
首先是数据结构需要支持遍历的能力,需要实现 Iterable 接口;另一部分重要的功能时添加和删除的方法。
定义优先级
class PriorityQueue<K, V>: Iterable<V> {
// 优先级的枚举定义
sealed class Priority(val level: Int) {
object High: Priority(3)
object Default: Priority(2)
object Low: Priority(1)
}
}
定义元素结构
因为选择了链表作为排序的数据结构,所以需要定义一下链表节点结构:
inner class Entry<V>(
val key: K? = null,
val priority: Priority = Priority.Default,
var next: Entry<V>? = null
)
底层数据结构
PriorityQueue 通过 HashMap 保存真实数据,通过链表来处理排序关系。
class PriorityQueue<K, V>: Iterable<V> {
// 排序链表的头节点
var head: Entry<V> = Entry()
// 三个优先级
private var lowHead: Entry<V>? = null
private var midHead: Entry<V>? = null
private var highHead: Entry<V>? = null
// 保存真实数据的底层数据结构
private val hashMap = HashMap<K, V>()
}
添加数据
添加数据类似 Map 的 put 方法,这里需要多一个优先级属性 Priority 参数,为了方便使用,这里默认设置了 Default 。
fun put(key: K?, value: V?, priority: Priority = Default): Boolean {
// 检查 key 和 value
if (key == null || value == null) return false
// 创建 链表节点对象
val newEntry: Entry<V> = Entry(key, priority)
// 加锁保证同步
synchronized(this) {
// 已存在包名,仅更新值
if (hashMap.containsKey(key)) {
hashMap[key] = value
return true
}
// 添加新节点
when (priority) {
Priority.Low -> {
// 插入到头节点
val temp = lowHead
lowHead = newEntry
newEntry.next = temp
}
Priority.Default -> {
val temp = midHead
midHead = newEntry
newEntry.next = temp
}
Priority.High -> {
val temp = highHead
highHead = newEntry
newEntry.next = temp
}
}
hashMap[key] = value
moveHead()
}
return false
}
在上面的逻辑中,在执行完添加逻辑后调用了 moveHead 方法,这个方法主要就是用来移动整体的排序链表的头节点:
private fun moveHead() {
head.next = highHead ?: midHead ?: lowHead
}
这个方法看似迷幻,实际上的含义很简单,按高中低优先级的链表依次取值,如果高优先级的头节点为空,取中优先级的头节点,中优先级也为空,取低优先级的头节点。
删除数据
1. 通过 key 删除
fun remove(key: K): Boolean {
// 先移除 HashMap 中的数据,移除的值赋值给 node
val node = hashMap.remove(key)
// 如果 node 存在,说明可以移除,在链表中查找并删除
if (node != null) {
var highTemp = highHead
var midTemp = midHead
var lowTemp = lowHead
var lastNode: Entry<V>? = null
// 查找,依次遍历 highTemp -> midTemp -> lowTemp
while (highTemp != null) {
if (highTemp.key == key) {
lastNode?.next = highTemp.next
return true
}
lastNode = highTemp
highTemp = highTemp.next
}
while (midTemp != null) {
if (midTemp.key == key) {
lastNode?.next = midTemp.next
return true
}
lastNode = midTemp
midTemp = midTemp.next
}
while (lowTemp != null) {
if (lowTemp.key == key) {
lastNode?.next = lowTemp.next
return true
}
lastNode = lowTemp
lowTemp = lowTemp.next
}
}
return false
}
2. 消费队列第一个元素
fun pop(): V? {
synchronized(this) {
val temp = head.next
// 根据 head 指向的第一个节点的优先级从不同的节点中删除即可
when (temp?.priority) {
Priority.High -> {
// head 节点为高优先级,处理高优先级的头节点和尾节点
val willRemoveNode = highHead
highHead = willRemoveNode?.next
}
Priority.Default -> {
val willRemoveNode = midHead
midHead = willRemoveNode?.next
}
Priority.Low -> {
val willRemoveNode = lowHead
lowHead = willRemoveNode?.next
}
}
moveHead()
return hashMap.remove(temp?.key)
}
}
迭代器和遍历实现
最后的部分时迭代器实现,因为 PriorityQueue 实现了 Iterable 接口,所以需要实现 fun iterator(): Iterator<V>
,迭代器 Iterator 需要进行自定义,Iterator 主要需要实现两个方法,hasNext() 和 next() :
inner class Itr: Iterator<V> {
var cursor: Entry<V>? = null
var highTemp = highHead
var midTemp = midHead
var lowTemp = lowHead
override fun hasNext(): Boolean {
// 检查排序链表是否为空
cursor = highTemp ?: midTemp ?: lowTemp
return cursor != null
}
override fun next(): V {
// 先找出排序链表的头节点
cursor = if (highTemp != null) {
highTemp
} else if (highTemp == null && midTemp != null) {
midTemp
} else if (midTemp == null && lowTemp != null) {
lowTemp
} else {
null
}
// 为空抛出 NoSuchElementException
if (cursor == null) {
throw NoSuchElementException("error")
}
// 根据头节点的优先级从不同的优先级链表中删除节点
when(cursor?.priority) {
Priority.High -> {
val willRemoveNode = highTemp
highTemp = willRemoveNode?.next
}
Priority.Default -> {
val willRemoveNode = midTemp
midTemp = willRemoveNode?.next
}
Priority.Low -> {
val willRemoveNode = lowTemp
lowTemp = willRemoveNode?.next
}
}
// 最后根据 entry 从 hashMap 中查找返回结果
return hashMap[cursor?.key] ?: throw NoSuchElementException("error")
}
}
然后实现 Iterable 接口中定义的 iterator 方法:
override fun iterator(): Iterator<V> {
return Itr()
}