JDK原始碼閱讀(十四) : 優先順序佇列——PriorityQueue
1. 優先順序佇列
PriorityQueue是一種基於堆的無界優先順序佇列。內部使用Object陣列儲存資料，在容量不足時會進行擴容操作。內部元素的排序規則，按照構造例項時傳入的Comparator或者元素自身的排序規則(所屬類實現Comparable介面)。
2. Fields
預設的陣列長度為11。
private static final int DEFAULT_INITIAL_CAPACITY = 11;
用於儲存資料的Object陣列。
transient Object[] queue;
元素數量。
private int 給元素排序的排序器，如果沒有自定義排序器，則會使用元素自身的排序規則。
private final Comparator<? super E> comparator;
容器修改次數。
transient int modCount = 0;
3. Constructors
（1）空參構造器。使用預設的陣列長度和元素自身的排序規則。
public PriorityQueue() {
this(DEFAULT_INITIAL_CAPACITY, null);
}
（2）自定義陣列長度的構造器。
public PriorityQueue(int initialCapacity) （3）自定義元素排序規則的構造器。
public PriorityQueue(Comparator<? super E> comparator) {
this(DEFAULT_INITIAL_CAPACITY, comparator);
}
（4）此外，還有其他構造器，可以自定義多種配置資訊。
4. 新增元素
add方法向優先順序佇列中新增元素。如果新增成功，會返回true。如果添加了null，會丟擲空指標異常。
public boolean add(E e) {
return offer public boolean offer(E e) {
//不可以新增null
if (e == null)
throw new NullPointerException();
modCount++;
//獲取當前元素數
int i = size;
//如果超出了陣列長度，則陣列擴容
if (i >= queue.length)
grow(i + 1);
//元素總數加1
size = i + 1;
//如果當前容器為空，則直接在0號位置放置元素
if (i == 0)
queue[0] = e;
else
//否則，將新元素插入到堆中，實行向上調整
siftUp(i, e);
return true;
}
【擴容策略】：如果舊陣列的長度小於64，則新陣列的長度變成舊陣列的2倍+2；如果舊陣列的長度大於等於64，則新陣列的長度為舊陣列的1.5倍。
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);
}
siftUp方法將新元素插入堆中時，根據是否使用自定義的排序器來執行不同的方法。
private void siftUp(int k, E x) {
if (comparator != null)
siftUpUsingComparator(k, x);
else
siftUpComparable(k, x);
}
siftUpComparable方法使用了元素自身的排序規則。前提是這個元素的所屬類必須已經實現Comparable介面。
private void siftUpComparable(int k, E x) {
//將新元素的型別轉換為Comparable
Comparable<? super E> key = (C<? super E>) x;
//起始時，k指向最後一個元素的後一位
//向上調整堆，直到到達堆頂
while (k > 0) {
//堆是一棵完全二叉樹
//取得最後一個元素的父親節點(父節點的座標為i，左子節點的座標為2*i+1，右子節點的座標為2*i+2)
//通過上述規則，如果某個節點的位置為k，則其父節點的位置為(k-1)/2
int parent = (k - 1) >>> 1;
//取得父節點元素
Object e = queue[parent];
//如果當前節點比父節點大，則可以直接插入在當前位置k上
if (key.compareTo((E) e) >= 0)
break;
//否則，就將父節點元素移動到當前位置k
queue[k] = e;
//k指向父節點
k = parent;
//在迴圈中，不斷向上調整，直到遇到當前節點比父節點大，或者，k到達了頂點，就能退出迴圈了
}
//最後，k為插入位置
queue[k] = key;
}
siftUpUsingComparator方法使用自定義排序器進行插入，與上述方法類似。
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;
}
5. 彈出堆頂元素
remove方法定義在PriorityQueue類的抽象父類AbstractQueue中，其內部也是呼叫了PriorityQueue類自身實現的poll方法。
在佇列為空時做彈出操作，remove方法會丟擲NoSuchElementException；而poll方法，會返回null。
public E remove() {
E x = poll();
if (x != null)
return x;
else
throw new NoSuchElementException();
}
public E poll() {
if (size == 0)
return null;
int s = --size;
modCount++;
//彈出的是堆頂元素
E result = (E) queue[0];
//取得最後一個元素
E x = (E) queue[s];
//將最後一個位置置為null，表示刪除
queue[s] = null;
if (s != 0)
//然後將前面取得的最後一個位置的元素放到堆頂，再向下調整堆
siftDown(0, x);
return result;
}
private void siftDown(int k, E x) {
if (comparator != null)
siftDownUsingComparator(k, x);
else
siftDownComparable(k, x);
}
siftDownComparable方法使用元素自身的排序規則向下調整堆。
private void siftDownComparable(int k, E x) {
//假設將最後一個元素key放在堆頂的位置，也就是k位置上
Comparable<? super E> key = (Comparable<? super E>)x;
//取得最後一個父節點的下一個位置
int half = size >>> 1; // loop while a non-leaf
//只需要調整到最後一個父節點，就可以結束。
while (k < half) {
//取得當前位置k的左孩子
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];
//比較k位置的元素和其最小的孩子，如果key更小，則可以停止調整了
if (key.compareTo((E) c) <= 0)
break;
//否則，就將更小的孩子放到k位置上
queue[k] = c;
//k下移到該孩子處
k = child;
//不斷調整，直到最後一個父節點也調整完畢，或者當前的父節點已經比兩個孩子都小了，就退出迴圈
}
//把key實際放入k位置上
queue[k] = key;
}
6. 優先順序佇列的應用
（1）求top K個元素。
在資料量很大的情況下，無法一次性將所有資料載入進記憶體進行排序來獲得topK個元素。此時，可以先將前面k個元素放入一個容量為k的優先順序佇列，然後一點一點把後面的資料載入到記憶體中，並嘗試插入堆中。
比如，我們要求最小K個元素，就建立一個大頂堆(也稱最大堆)，堆頂元素是最大的。所以排序規則需要按照從大到小排序。
插入的時候，首先比較堆頂元素，如果大於等於堆頂，就不用插入了；如果小於堆頂，就將堆頂移除，然後將新元素插入堆中。
使用這種方式，可以動態地求得top K個元素。
//numbers陣列儲存了全部資料
public void topK(int[] numbers, int k){
int len;
if(numbers == null || (len = numbers.length) == 0 || k > len){
throw new IllegalArgumentException();
}
//建立一個長度為k的大頂堆
Comparator<Integer> comparator = (o1, o2) -> {
if (o1.compareTo(o2) == -1) {
return 1;
} else if (o1.equals(o2)) {
return 0;
} else {
return -1;
}
};
//優先順序佇列就是一個堆，傳入自定義的排序器，從大到小排序
PriorityQueue<Integer> pq = new PriorityQueue<>(k,comparator);
//將前k個元素插入堆中，時間複雜度為O(klogk)
for(int i = 0; i < k; ++i){
pq.add(numbers[i]);
}
//將剩餘元素插入到堆中，時間複雜度O(nlogk)
for(int i = k; i < len; ++i){
//取得堆頂元素
Integer peek = pq.peek();
//如果當前元素比堆頂元素小，就將堆頂元素移除，將當前元素插入
if(numbers[i] < peek){
pq.remove();
pq.add(numbers[i]);
}
}
//輸出堆中的元素
for(int i : pq){
System.out.print(i + " ");
}
}
（2）求動態資料的中位數。
可以用兩個堆來求中位數。資料總數為n。
如果n為偶數，則將前n/2個元素放入到一個大頂堆中，將後n/2個元素放入到一個小頂堆中。此時，兩個堆的堆頂是中位數。
如果n為奇數，則將前(n+1)/2個元素放入到一個大頂堆中，將後(n-1)/2個元素放入到一個小頂堆中。此時，大頂堆的堆頂是中位數。
//建兩個堆，一個是大頂堆，一個是小頂堆
private PriorityQueue<Integer> maxpq = new PriorityQueue<>((o1, o2) -> {
if (o1.compareTo(o2) == -1) {
return 1;
} else if (o1.equals(o2)) {
return 0;
} else {
return -1;
}
});
private PriorityQueue<Integer> minpq = new PriorityQueue<>();
//numbers陣列用於儲存資料
private int[] numbers;
//size儲存資料總數
private int size;
將已有的資料放入到兩個堆中。
public void initHeaps(){
//將前一半資料放入大頂堆
int mid = (size - 1) >>> 1;
for(int i = 0; i <= mid; ++i){
maxpq.add(numbers[i]);
}
//遍歷後一半資料，如果比大頂堆的堆頂元素小，就將這個堆頂元素移動到小頂堆，並將當前資料放入大頂堆
for(int i = mid+1; i < size; ++i){
Integer peek = maxpq.peek();
if(peek > numbers[i]){
Integer move = maxpq.remove();
maxpq.add(numbers[i]);
minpq.add(move);
}else {//否則，就將當前資料直接放入小頂堆
minpq.add(numbers[i]);
}
}
}
如果陣列中存在動態新增元素的操作，中位數不斷地變化，那麼在新增元素時，如果新元素小於大頂堆的堆頂元素，則將其插入到大頂堆；如果大於等於大頂堆的堆頂元素，則插入到小頂堆。但是，這種做法可能會導致兩個堆不平衡，此時只需要在每次插入操作後，對兩個堆進行調整，將數量多的堆的堆頂元素不斷移動到數量少的堆，直到兩個堆平衡。時間複雜度為O(logn)。
public void add(int e){
ensureCapacity(size+1);
numbers[size++] = e;
//插入到兩個堆中
Integer peek = maxpq.peek();
if(peek == null) {
maxpq.add(e);
}else {
if (e < peek) {
maxpq.add(e);
} else {
minpq.add(e);
}
}
//調整兩個堆的大小
balanceTwoHeaps();
}
調整兩個堆的大小。始終保持大頂堆比小頂堆多一個，或者兩個堆數量相等。
private void balanceTwoHeaps() {
int size1 = maxpq.size();
int size2 = minpq.size();
while(!(size1 == size2) && !(size1 == size2 + 1)){
if(size1 > size2){
minpq.add(maxpq.remove());
--size1;
++size2;
}else {
maxpq.add(minpq.remove());
++size1;
--size2;
}
}
}
獲取中位數。如果資料總數是奇數，則大頂堆的堆頂元素為中位數。如果是偶數，則兩個堆頂的元素為中位數。時間複雜度為O(1)。
public void getMiddle(){
if((size&1) == 1){
System.out.println(maxpq.peek());
}else {
System.out.println(maxpq.peek() + " " + minpq.peek());
}
}
【舉一反三】：取中位數可以用兩個堆來實現，取其他位置的元素同樣可以用兩個堆來實現。