走進 JDK 之 PriorityQueue

走進 JDK 系列第 16 篇

文章相關源碼： PriorityQueue.java

這是 Java 集合框架的第三篇文章了，前兩篇分別解析了 ArrayList 和 LinkedList，它們分別是基於動態數組和鏈表來實現的。今天來說說 Java 中的優先級隊列 PriorityQueue，它是基於堆實現的，後面也會介紹堆的相關概念。

概述

PriorityQueue 是基於堆實現的無界優先級隊列。優先級隊列中的元素順序根據元素的自然序或者構造器中提供的 Comparator。不允許 null 元素，不允許插入不可比較的元素(未實現 Comparable)。它不保證線程安全，JDK 也提供了線程安全的優先級隊列 PriorityBlockingQueue。

劃個重點，基於堆實現的優先級隊列。首先來看一下什麼是隊列？什麼是堆？

隊列

隊列其實很好理解，它是一種特殊的線性表。比如食堂排隊打飯就是一個隊列，先排隊的人先打飯，後來的同學在隊尾排隊，打到飯的同學從對頭離開，這就是典型的 先進先出（FIFO）隊列。隊列一般會提供 入隊 和 出隊 兩個基本操作，入隊在隊尾進行，出隊在對頭進行。Java 中隊列的父類接口是 Queue。我們來看一下 Queue 的 uml 圖，給我們提供了哪些基本方法：

• add(E) : 在隊尾添加元素
• offer(E) : 在隊尾添加元素。在容量受限的隊列中和 add() 表現一致。
• remove() : 刪除並返回隊列頭，隊列爲空時拋出異常
• poll() : 刪除並返回隊列頭，隊列爲空時返回 null
• element（）： 返回隊列頭，但不刪除，隊列爲空時拋出異常
• peek() : 返回隊列頭，但不刪除，隊列爲空時返回 null

基本也就是對出隊和入隊操作進行了細分。PriorityQueue 是一個優先級隊列，會按自然序或者提供的 Comparator 對元素進行排序，這裏使用的是堆排序，所以優先級隊列是基於堆來實現的。如果你瞭解堆的概念，就可以跳過下一節了。如果你不知道什麼是堆，仔細閱讀下一節，不然是沒辦法理解 PriorityQueue 的源碼的。

堆

堆其實是一種特殊的二叉樹，它具備如下兩個特徵：

• 堆是一個完全二叉樹
• 堆中每個節點的值都必須小於等於（或者大於等於）其子節點的值

對於一個高度爲 k 的二叉樹，如果它的 0 到 k-1 層都是滿的，且最後一層的所有子節點都是在左邊那麼他就是完全二叉樹。用數組實現的完全二叉樹可以很方便的根據父節點的下標獲取它的兩個子節點。下圖就是一個完全二叉樹：

堆就是一個完全二叉樹。頂部是最小元素的叫小頂堆，頂部是最大元素的叫大頂堆。PriorityQueue 是小頂堆。對照上面的堆結構，對於任意父節點，以下標爲 4 的節點 5 爲例，它的兩個子節點下標分別爲 2*4+1 和 2*4+2。關於完全二叉樹和堆，記住下面幾個結論，都是後面的源碼分析中要用到的：

• 沒有子節點的節點叫做葉子節點
• 下標爲 n 的父節點的兩個左右子節點的下標分別是 2*n+1 和 2*n+2

這就是用數組來構建堆的好處，根據下標就可以快速構建堆結構。堆就先說到這裏，記住優先級隊列 PriorityQueue 是基於堆實現的隊列，堆是一個完全二叉樹。下面就根據 PriorityQueue 的源碼對堆的操作進行深入解析。

源碼解析

類聲明

public class PriorityQueue<E> extends AbstractQueue<E>
    implements java.io.Serializable { }

成員變量

private static final long serialVersionUID = -7720805057305804111L;
private static final int DEFAULT_INITIAL_CAPACITY = 11; // 默認初始容量
transient Object[] queue; // 存儲隊列元素的數組
private int size = 0; // 隊列元素個數
private final Comparator<? super E> comparator; 
transient int modCount = 0; // fail-fast
private static final int MAX_ARRAY_SIZE = Integer.MAX_VALUE - 8; // 最大容量

• PriorityQueue 使用數組 queue 來存儲元素，默認初始容量是 11，最大容量是 Integer.MAX_VALUE - 8。
• comparator 若爲 null，則按照元素的自然序來排列。
• modCount 用來提供 fail-fast 機制。

構造函數

PriorityQueue 的構造函數有 7 個，可以分爲兩類，提供初始元素和不提供初始元素。先來看看不提供初始元素的構造函數：

/*
 * 創建初始容量爲 11 的優先級隊列，元素按照自然序
 */
public PriorityQueue() {
    this(DEFAULT_INITIAL_CAPACITY, null);
}

/*
 * 創建指定初始容量的優先級隊列，元素按照自然序
 */
public PriorityQueue(int initialCapacity) {
    this(initialCapacity, null);
}

/*
 * 創建初始容量爲 11 的優先級隊列，元素按照按照給定 comparator 排序
 */
public PriorityQueue(Comparator<? super E> comparator) {
    this(DEFAULT_INITIAL_CAPACITY, comparator);
}

/*
 * 創建指定初始容量的優先級隊列，元素按照按照給定 comparator 排序
 */
public PriorityQueue(int initialCapacity,
                        Comparator<? super E> comparator) {
    if (initialCapacity < 1)
        throw new IllegalArgumentException();
    this.queue = new Object[initialCapacity];
    this.comparator = comparator;
}

這一類構造函數都很簡單，直接給 queue 和 comparator 賦值即可。對於給定初始元素的構造函數就沒有這麼簡單了，因爲給定的初始集合並不一定滿足堆的結構，我們需要將其構造成堆，這個過程稱之爲 堆化。

PriorityQueue 可以直接根據 SortedSet 和 PriorityQueue 來構造堆，由於初始集合本來就是有序的，所以無需進行堆化。如果構造器參數是任意 Collection，那麼就可能需要堆化了。

public PriorityQueue(Collection<? extends E> c) {
    if (c instanceof SortedSet<?>) { // 直接使用 SortedSet 的 comparator
        SortedSet<? extends E> ss = (SortedSet<? extends E>) c;
        this.comparator = (Comparator<? super E>) ss.comparator();
        initElementsFromCollection(ss);
    }
    else if (c instanceof PriorityQueue<?>) { // 直接使用 PriorityQueue 的 comparator
        PriorityQueue<? extends E> pq = (PriorityQueue<? extends E>) c;
        this.comparator = (Comparator<? super E>) pq.comparator();
        initFromPriorityQueue(pq);
    }
    else {
        this.comparator = null;
        initFromCollection(c); // 需要堆化
    }
}

我們來看看堆化的具體過程：

private void initFromCollection(Collection<? extends E> c) {
    initElementsFromCollection(c); // 將 c copy 一份直接賦給 queue
    heapify(); // 堆化
}

private void heapify() {
    for (int i = (size >>> 1) - 1; i >= 0; i--)
        siftDown(i, (E) queue[i]); // 自上而下堆化
}

堆化的邏輯很短，但是內容很豐富。堆化其實用兩種，shiftDown() 是自上而下堆化，shiftUp() 是自下而上堆化。這裏使用的是 shiftDown。從上面的代碼中你可以看出從哪一個結點開始堆化的嗎？並不是從最後一個節點開始堆化，而是從最後一個非葉子節點開始的。還記得什麼是葉子節點嗎，沒有子節點的節點就是葉子節點。所以，對所有非葉子節點進行堆化，就足以處理所有節點了。那麼最後一個非葉子節點的下標是多少呢，如果想不出來可以翻到上面的堆的示意圖，答案就是 size/2 - 1，源碼中使用了無符號移位操作代替了除法。

再來看看 shiftDown() 的具體邏輯：

/*
 * 自上而下堆化，保證 x 小於等於子節點或者 x 是一個葉子結點
 */
private void siftDown(int k, E x) {
    if (comparator != null)
        siftDownUsingComparator(k, x);
    else
        siftDownComparable(k, x);
}

x 是要插入的元素，k 是要填充的位置。根據 comparator 是否爲空調用不同的方法。這裏以 comparator 不爲 null 爲例：

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) { // 堆的非葉子節點的個數總是小於 half 的。當 k 是葉子節點的時候，直接交換即可
        int child = (k << 1) + 1; // 左子節點
        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) // 比較 x 和子節點
            break; // x 比子節點大，直接跳出循環
        queue[k] = c; // 若 x 比子節點小，和子節點交換
        k = child; // 此時 k 等於 child，繼續和子節點比較
    }
    queue[k] = key;
}

邏輯比較簡單。PriorityQueue 是一個小頂堆，父節點總是小於等於子節點。對於每一個非葉子節點，將它和自己的兩個左右子節點進行比較，若父節點比兩個子節點都大，就要將這個父節點下沉，下沉之後再繼續和子節點比較，直到該父節點比兩個子節點都小，或者這個父節點已經是葉子結點，沒有子節點了。這樣循環往復，自上而下的堆化就完成了。

方法

看完了構造函數，我們來看看 PriorityQueue 提供的方法。既然是隊列，那就肯定有入隊和出隊操作。先來看看入隊方法 add()。

add(E e)

public boolean add(E e) {
    return offer(e);
}

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;
}

add() 方法會調用 offer() 方法，它們都是在隊尾增加一個元素。offer() 過程可以分爲兩步：自動擴容 和 堆化 。

優先隊列也支持自動擴容，但其擴容邏輯和 ArrayList 不同，ArrayList 是直接擴容至原來的 1.5 倍。而 PriorityQueue 根據當前隊列大小的不同有不同的表現。

private void grow(int minCapacity) {
    int oldCapacity = queue.length;
    // Double size if small; else grow by 50%
    int newCapacity = oldCapacity + ((oldCapacity < 64) ?
                                     (oldCapacity + 2) :
                                     (oldCapacity >> 1));
    // overflow-conscious code
    if (newCapacity - MAX_ARRAY_SIZE > 0)
        newCapacity = hugeCapacity(minCapacity);
    queue = Arrays.copyOf(queue, newCapacity);
}

• 原隊列大小小於 64 時，直接翻倍再 +2
• 原隊列大小大於 64 時，增加 50%

第二步就是堆化了。隊尾增加元素爲什麼要重新堆化呢？看下面這個圖：

左邊是一個堆，我要在隊尾添加一個元素 4，如果這樣直接加在隊尾，還是一堆嗎？顯然不是的了，因爲 4 比 5 小，卻排在了 5 的下面。所以這時候就需要堆化了。前面介紹過 shiftDown, 這裏還可以自上而下堆化嗎？顯然不行，因爲在隊尾添加節點，這個節點肯定是葉子節點，它已經位於最下面一層了，沒有子節點了。這就要使用另一種堆化方法，自下而上堆化。拿 4 和其父節點比較，發現 4 比 5 小，和父節點交換，這時候 4 就處在 下標爲 2 的位置了。再和父節點比較，發現 4 比 1 大，不交換，結束堆化。這時候 4 就找到自己在堆中正確的位置了。對應源代碼中的 shiftUp() 方法：

private void siftUp(int k, E x) {
    if (comparator != null)
        siftUpUsingComparator(k, x);
    else
        siftUpComparable(k, x);
}
    
private void siftUpUsingComparator(int k, E x) {
    while (k > 0) {
        int parent = (k - 1) >>> 1; // 找到 k 位置的父節點
        Object e = queue[parent];
        if (comparator.compare(x, (E) e) >= 0) // 比較 x 與父節點值的大小
            break; // x 比父節點大，直接跳出循環
        queue[k] = e; // 若 x 比父節點小，交換元素
        k = parent; // 此時 k 等於 parent，繼續和父節點比較
    }
    queue[k] = x;
}

根據下標 k 就可以找到 k 位置的父節點，這也是前面介紹堆的時候給出的結論。那麼其插入操作的時間複雜度是多少呢？這和堆的高度相關，最好時間複雜度就是 O(1)，不需要交換元素，最壞時間複雜度是 O(log(n))，因爲堆的高度是 log(n)，最壞情況就是一路交換到堆頂。平均時間複雜度也就是 O(log(n))。

說完了入隊，下面看一下出隊。

poll()

poll() 是出隊操作，也就是移除隊隊頭元素。想想一下，一個完全二叉樹，你把堆頂移除了，它就不是一個完全二叉樹了，也就沒辦法去堆化了。源碼中是這樣處理的，移除隊頭元素之後，暫時把隊尾元素移到隊頭，這樣它又是一個完全二叉樹了，就可以進行堆化了。下面這個圖更容易理解：

這裏的堆化操作很顯然應該是 shiftDown() 了，自上而下堆化。

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;
}

除了移除隊列頭，PriorityQueue 也支持 remove 任意位置的節點，通過 remove() 方法實現。

remove()

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]; // moved 是隊尾元素
        queue[s] = null;
        siftDown(i, moved); // 將隊尾元素插入 i 位置，再自上而下堆化
        if (queue[i] == moved) {
            siftUp(i, moved); // moved 沒有往下交換，仍然在 i 位置處，此時需要再自下而上堆化以保證堆的正確性
            if (queue[i] != moved)
                return moved;
        }
    }
    return null;
}

如果是刪除隊尾，直接刪除皆可以了。但如果是刪除中間某個節點，就會在堆中形成一個空洞，不再是完全二叉樹。其實和 poll 的處理方式一致，將隊尾節點暫時填充到刪除的位置，形成完全二叉樹再進行堆化。

這裏的堆化過程和 poll 有一些不一致。首先進行 shiftDown()，自上而下堆化。shiftDown() 完成之後比較 queue[i] == moved，如果不相等，說明節點 i 向下交換了，它找到了自己的位置。但是如果相等，則說明節點 i 沒有向下交換，也就是節點 i 的值比它的子節點都要小。但這並不能說明它一定比它的父節點大。所以，這種情況還需要再自下而上堆化，以保證可以完全符合堆的特性。

總結

說了半天 PriorityQueue ，其實都是在說堆。如果你對堆很熟悉的話，PriorityQueue 的源碼很好理解。當然不熟悉也沒關係，藉着源碼正好可以學習一下堆的基本概念。最後簡單總結一下優先隊列：

• PriorityQueue 是基於堆的，堆是一個特殊的完全二叉樹，它的每一個節點的值都必須小於等於（或大於等於）其子樹中每個節點的值
• PriorityQueue 的出隊和入隊操作時間複雜度都是 O(log(n))，僅與堆的高度有關
• PriorityQueue 初始容量爲 11，支持動態擴容。容量小於 64 時，擴容一倍。大於 64 時，擴容 0.5 倍
• PriorityQueue 不允許 null 元素，不允許不可比較的元素
• PriorityQueue 不是線程安全的，PriorityBlockingQueue 是線程安全的

PriorityQueue 就說到這裏了。下一篇應該會寫 Set 相關。

文章首發微信公衆號： 秉心說 ， 專注 Java 、 Android 原創知識分享，LeetCode 題解。

更多 JDK 源碼解析，掃碼關注我吧！