走进 JDK 之 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; // 最大容量
复制代码

/*
 * 创建初始容量为 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;
}
复制代码

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]); // 自上而下堆化
}
复制代码

/*
 * 自上而下堆化，保证 x 小于等于子节点或者 x 是一个叶子结点
 */
private void siftDown(int k, E x) {
    if (comparator != null)
        siftDownUsingComparator(k, x);
    else
        siftDownComparable(k, x);
}
复制代码

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;
}
复制代码

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;
}
复制代码

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);
}
复制代码

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;
}
复制代码

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;
}
复制代码

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;
}
复制代码