PriorityQueue原理分析——基于源码

Posted 2020-12-21 buptleida

在业务场景中，处理一个任务队列，可能需要依照某种优先级顺序，这时，Java中的PriorityQueue（优先队列）便可以派上用场。优先队列的原理与堆排序密不可分，可以参考我之前的一篇博客：

堆排序总结与实现

原理

PriorityQueue中维护一个Queue[]数组，在逻辑上把它理解成一个小根堆或大根堆，即一个完全二叉树，每一个三元组中父节点小于两个孩子结点（小根堆，如果是大于则是大根堆）。本博客以小根堆来进行说明，因为PriorityQueue默认实现小根堆，即小的数先出队，当然也可以自定义Comparator实现大根堆。

• 入队：每次入队时，把新元素挂在最后，从下往上遍历调整成小根堆；
• 出队：每次出队时，移除顶部元素，把最后的元素移到顶部，并从上往下遍历调整成小根堆。

出队

poll()方法如下：

public E poll() {
    if (size == 0)
        return null;
    int s = --size;
    modCount++;
    E result = (E) queue[0];
    E x = (E) queue[s];
    queue[s] = null;
    if (s != 0)
        siftDown(0, x);
    return result;
}

可以看到，队首元素 queue[0] 出队，队尾的元素 queue[s] 进入 siftDown(0, x) 方法进行堆调整。siftDown方法如下：

private void siftDown(int k, E x) {
    if (comparator != null)
        siftDownUsingComparator(k, x);
    else
        siftDownComparable(k, x);
}
//k为开始遍历的位置，x为需要插入的值
@SuppressWarnings("unchecked")
private void siftDownComparable(int k, E x) {
    Comparable<? super E> key = (Comparable<? super E>)x;
    int half = size >>> 1;        // loop while a non-leaf
    // 只需要遍历到数组的一半即可，保证遍历到最后一个三元组的父节点即可
    while (k < half) {
        int child = (k << 1) + 1; // assume left child is least
        Object c = queue[child];
        int right = child + 1;
        if (right < size &&
            ((Comparable<? super E>) c).compareTo((E) queue[right]) > 0)
            c = queue[child = right];//比较左右孩子结点，取最小的那个
        if (key.compareTo((E) c) <= 0)
            break;//找到了key应该放入的位置
        queue[k] = c;
        k = child;
    }
    queue[k] = key;
}

@SuppressWarnings("unchecked")
private void siftDownUsingComparator(int k, E x) {
    int half = size >>> 1;
    while (k < half) {
        int child = (k << 1) + 1;
        Object c = queue[child];
        int right = child + 1;
        if (right < size &&
            comparator.compare((E) c, (E) queue[right]) > 0)
            c = queue[child = right];
        if (comparator.compare(x, (E) c) <= 0)
            break;
        queue[k] = c;
        k = child;
    }
    queue[k] = x;
}

可以看到，这与堆排序中的堆调整如出一辙。

入队

offer方法如下所示：

public boolean offer(E e) {
    if (e == null)
        throw new NullPointerException();
    modCount++;
    int i = size;
    if (i >= queue.length)
        grow(i + 1);
    size = i + 1;
    if (i == 0)
        queue[0] = e;
    else
        siftUp(i, e);
    return true;
}

同样，其核心在于 siftUp(i, e) 方法。如下所示：

private void siftUp(int k, E x) {
    if (comparator != null)
        siftUpUsingComparator(k, x);
    else
        siftUpComparable(k, x);
}

@SuppressWarnings("unchecked")
private void siftUpComparable(int k, E x) {
    Comparable<? super E> key = (Comparable<? super E>) x;
    while (k > 0) {
        int parent = (k - 1) >>> 1;//结点父节点的下标
        Object e = queue[parent];
        if (key.compareTo((E) e) >= 0)
            break;//如果结点值大于父节点，则可以放置在该三元组下
        queue[k] = e;//向子节点赋值父节点的值，不用担心某些值被覆盖，因为初始k等于size
        k = parent;
    }
    queue[k] = key;//最后在待插入位置赋key的值
}

@SuppressWarnings("unchecked")
private void siftUpUsingComparator(int k, E x) {
    while (k > 0) {
        int parent = (k - 1) >>> 1;
        Object e = queue[parent];
        if (comparator.compare(x, (E) e) >= 0)
            break;
        queue[k] = e;
        k = parent;
    }
    queue[k] = x;
}

此方法，是一个不断从父节点往子节点赋值的过程，直到找到适合放置插入结点值的位置。

移除

removeAt 方法如下所示：

private E removeAt(int i) {
    // assert i >= 0 && i < size;
    modCount++;
    int s = --size;
    if (s == i) // removed last element
        queue[i] = null;
    else {
        E moved = (E) queue[s];
        queue[s] = null;
        siftDown(i, moved);
        if (queue[i] == moved) {
            siftUp(i, moved);
            if (queue[i] != moved)
                return moved;
        }
    }
    return null;
}

移除下标为i的元素，相当于以 i 为根节点的完全二叉树的出队，于是执行 siftDown 方法调整最后一个元素 moved 的位置，即将该堆调整为小根堆。调整完之后，如果 moved 没有来到 i 的位置，说明 i 以上的堆结构一定符合规则；如果 moved 被调整到 i 位置，i上面的父节点有可能比 moved大，所以需要 siftUp(i, moved) 方法从 i 位置向上调整，调整为小根堆，完毕。

总结

其实不管是 siftUp 方法还是 siftDown 方法，都是利用了完全二叉树的性质，通过父节点与孩子结点之间的快速访问来实现的。