PriorityQueue原理分析——基於原始碼

數小錢錢的種花兔發表於2020-11-13
原始碼

在業務場景中，處理一個任務佇列，可能需要依照某種優先順序順序，這時，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 方法，都是利用了完全二叉樹的性質，通過父節點與孩子結點之間的快速訪問來實現的。

相關文章