实现一个后入先出的优先级队列

63 阅读5分钟

实现一个后入先出的优先级队列

最近需要做个优先级队列来做任务,虽然 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()
    }