使用Kotlin实现Java的优先队列PriorityQueue

前言

上周在面试时，偶然一个算法用到了优先队列思想。我只懂效果不懂实现，当时感觉和堆排序的思想差不多。今天深入源码，自己又实现一遍加深印象。

源码有什么

1. 具有Queue和Collection集合和Queue队列的性质
2. 可以保证每次取出的元素都是最值（默认是最小，可以自己设置）
3. 内部采用推排序思想，上浮siftUp和下沉siftDown
4. 存储采用可变数组（和ArrayList一样），默认大小是11，刚开始每次*2+2，后面每次加一半（size+=size>>2）

如何实现一个优先队列

1. 队列最基本的offer，peek，poll
2. 集合最基本的isEmpty，toString，size
3. 供内部使用的siftUp，siftDown，grow

代码实现

个人有写算法加注释的习惯，看不懂私聊

import java.util.*
import kotlin.math.max

fun main() {
    val queue = MyPriorityQueue().apply {
        repeat(10) {
            offer(Random().nextInt(100))
        }
    }
    while (queue.isNotEmpty()) {
        println(queue.poll())
    }
}

class MyPriorityQueue {
    // 保存元素，方便实现泛型，默认大小是11
    private var array = Array(11) { 0 }

    // 当前数目
    private var size = 0

    // 增
    fun offer(value: Int) {
        if (size == array.size) {
            grow(size + 1)
        }
        array[size++] = value
        siftUp(size - 1, value)
    }


    // 看队头
    fun peek(): Int {
        return if (size == 0) -1 else array[0]
    }

    // 取对头
    fun poll(): Int {
        if (size == 0) {
            return -1
        }
        val result = array[0]
        array[0] = array[size - 1]
        size--
        siftDown(0, array[0])
        return array[0]
    }

    // 没有size方法
    fun size() = size

    // 没有isEmpty方法
    fun isEmpty() = size == 0

    fun isNotEmpty() = size > 0

    // 上浮，默认最小在上面
    private fun siftUp(index: Int, value: Int) {
        // 在上面选个大的和当前交换，没有退出
        var i = index
        while (i > 0) {
            val parent = (i - 1) / 2
            if (array[parent] <= value) {
                break
            }
            array[i] = array[parent]
            i = parent
        }
        array[i] = value
    }

    // 下沉，默认最大在下面
    private fun siftDown(index: Int, value: Int) {
        // 选出子节点最小的一个，没有则退出
        var i = index
        // 当至少存在左子节点时
        while (i * 2 + 1 < size) {
            // 找出最小的子节点下标
            val minChild = if (i * 2 + 2 >= size || array[i * 2 + 1] <= array[i * 2 + 2]) i * 2 + 1 else i * 2 + 2
            // 如果最小的还是比父节点大，退出
            if (array[minChild] >= value) {
                break
            }
            array[i] = array[minChild]
            i = minChild
        }
        array[i] = value
    }

    // 扩容机制，小于64，*2+2，大于，+>>2
    private fun grow(minCapacity: Int) {
        var newCapacity = array.size + if (array.size < 64) array.size + 2 else array.size / 2
        // 在一次添加大量元素才可能用到
        newCapacity = max(newCapacity, minCapacity)
        array = Arrays.copyOf(array, newCapacity)
    }

    // 重写toString方法
    override fun toString() = StringBuilder().apply {
        append("[")
        for (i in 0 until size) {
            append(
                if (i == 0) array[i] else " ${array[i]}"
            )
        }
        append("]")
    }.toString()
}