Java8 - PriorityQueue源码

一、概述

先来看看源码中的第一段注释，我们先尝试从中提取一些信息：

An unbounded priority queue based on a priority heap.The elements of the priority queue are ordered according to their Comparable natural ordering, or by a Comparator provided at queue construction time, depending on which constructor isused. A priority queue does not permit null elements.A priority queue relying on natural ordering also does not permit insertion of non-comparable objects (doing so may result in ClassCastException).

这段注释说的是：这是一个基于堆的无界优先队列，优先队列里的所有元素都根据它们的自然顺序排序，或者根据创建时提供的比较器来排序。优先队列里不允许 null 元素。另外，如果是根据自然顺序排序的优先队列，是不允许插入不可比较的对象，否则可能会抛出 ClassCastException。

从上面我们可以了解到，PriorityQueue 是一个基于堆排序实现的，带有优先级的队列，而这个优先级是指自然顺序或者是自定义的优先级（通过比较器），每次我们从队列里获取队头元素的时候，它都会给我们优先级最大（最小）的元素。

二、堆排序

想要了解 PriorityQueue 的实现原理，很有必要了解一下堆排序的相关概念：

这里的堆本质上是指一棵完全二叉树。如下图：

同时，根据排序优先级的不同，这个堆分为小顶堆和大顶堆。小顶堆是指父节点的值小于等于任何一个子节点的值，此时堆顶元素是堆中元素的最小值，如上图所示；反之则称为大顶堆，它的堆顶元素是堆中元素的最大值。

另外，虽然这个堆在逻辑结构上是一个树型结构，但是基于完全二叉树的特殊性，在实际中，我们通常使用数组来作为存储结构。如下图：

从上图我们可以看出，节点 i 的左儿子为 2*i+1，右儿子为 2*i+2；相反地，节点 i 的父亲节点为 (i-1)/2。通过这种方法，即使是使用数组作为存储结构，我们也可以快速定位到一个节点的父亲节点或子节点，完全不慌。

三、属性

简单了解了堆排序之后，我们重新进入正题，先来看一看 PriorityQueue 中定义了哪些属性：

//默认容量
private static final int DEFAULT_INITIAL_CAPACITY = 11;

//底层的存储结构，也是通过数组来存储的
transient Object[] queue; // non-private to simplify nested class access

//队列中的元素个数
private int size = 0;

//比较器，默认为null
private final Comparator<? super E> comparator;

//queue的最大容量
private static final int MAX_ARRAY_SIZE = Integer.MAX_VALUE - 8;

可以发现，PriorityQueue 的底层实现也是通过一个数组来作为存储结构，数组的默认大小是11。另外，它还有一个比较器，默认情况下比较器为 null。当比较器为空时，它会按照元素的自然顺序进行排序，如果元素无法比较，则会抛出异常。

四、方法

1、构造方法

PriorityQueue 提供了很多构造方法，我们可以根据自己的需求选择合适的构造方法来使用，具体如下：

//无参的构造方法，默认容量大小为11，比较器为null
public PriorityQueue() {
    this(DEFAULT_INITIAL_CAPACITY, null);
}

//指定初始容量大小，比较器为null
public PriorityQueue(int initialCapacity) {
    this(initialCapacity, null);
}

//指定比较器，初始容量大小为11
public PriorityQueue(Comparator<? super E> comparator) {
    this(DEFAULT_INITIAL_CAPACITY, comparator);
}

//指定初始容量大小和比较器
public PriorityQueue(int initialCapacity, Comparator<? super E> comparator) {
    // Note: This restriction of at least one is not actually needed,
    // but continues for 1.5 compatibility
    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<? extends E> ss = (SortedSet<? extends E>) c;
        this.comparator = (Comparator<? super E>) ss.comparator();
        initElementsFromCollection(ss);
    }
    else if (c instanceof PriorityQueue<?>) {
        PriorityQueue<? extends E> pq = (PriorityQueue<? extends E>) c;
        this.comparator = (Comparator<? super E>) pq.comparator();
        initFromPriorityQueue(pq);
    }
    else {
        this.comparator = null;
        initFromCollection(c);
    }
}

//通过一个优先队列来初始化
public PriorityQueue(PriorityQueue<? extends E> c) {
    this.comparator = (Comparator<? super E>) c.comparator();
    initFromPriorityQueue(c);
}

//通过一个SortedSet来初始化优先队列
public PriorityQueue(SortedSet<? extends E> c) {
    this.comparator = (Comparator<? super E>) c.comparator();
    initElementsFromCollection(c);
}

---

2、扩容

在上面，我们知道了 PriorityQueue 的底层是通过数组来存储的，并且数组的默认大小为11，所以对于 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);
}

private static int hugeCapacity(int minCapacity) {
    if (minCapacity < 0) // overflow
        throw new OutOfMemoryError();
    return (minCapacity > MAX_ARRAY_SIZE) ?
        Integer.MAX_VALUE :
        MAX_ARRAY_SIZE;
}

从上面的源码中，我们可以看出，当 queue 的长度小于 64 时，它每次扩容都是比原来大一倍且再多扩充两个存储单位；而当 queue 的长度大于等于 64 之后，它每次扩容都是比原来多扩50%的空间。

另外，其实在扩容的时候，newCapacity 的计算是会出现溢出的情况的，但是代码中通过 “newCapacity - MAX_ARRAY_SIZE > 0” 这一句来处理溢出的情况了（想一想这样为什么能处理）。如果这一句写成了 “newCapacity > MAX_ARRAY_SIZE” 那就 GG 了。

---

3、添加元素

在往 PriorityQueue 里添加元素的时候，每次都默认添加到队尾去，然后一直与父节点进行比较，判断是否需要调整节点的值，直至使它满足大顶堆（或小顶堆）的性质。这个过程通过图片来描述会比较容易理解一些，例如下面演示了往一个堆中添加一个元素 9 的过程，如图所示：

一开始，9 被添加到最后面，此时破坏了这个小顶堆的性质，所以接下来会开始做调整。先拿 9 跟其父节点比较，发现父节点 13 比较大，所以交换其位置；之后继续跟父节点 11 比较，发现父节点 11 比较大，继续交换位置；继续再跟父节点 8 比较，发现 9 比较大，不必交换位置，调整结束。此时这个堆又满足了小顶堆的性质，即父节点大于等于任何一个子节点。

很明显，在向上调整的过程中，最多只会调整 log(n) 个节点，所以这个操作的复杂度为 O(logn)。

添加元素的源码如下，结合这张图片，源码还是比较好理解的。

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

public boolean offer(E e) {
    //不允许null元素
    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 siftUp(int k, E x) {
    if (comparator != null)
        siftUpUsingComparator(k, x);
    else
        siftUpComparable(k, x);
}

//通过自然顺序来比较
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 = parent;
    }
    queue[k] = key;
}

//通过比较器来比较
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;
}

---

4、删除元素

PriorityQueue 删除元素有两种方法，一种是获取并删除堆顶元素，另外一种是删除指定元素，我们先来看一下获取并删除堆顶元素的过程：

获取堆顶元素很简单，直接取 queue[0] 就可以了，但是删除了堆顶元素之后，此时就需要对堆的结构进行调整，重新选举出一个堆顶。从上面的图片，我们可以看到，实际上堆顶并没有被删除，而是在调整的过程中被覆盖掉了，被删除的是堆中的最后一个元素 13。堆的调整过程其实是在为 13 这个元素寻找一个适合的位置插入，保证在 13 插入后，这棵树仍然满足堆的性质。调整的策略很简单，例如图片中的小顶堆，它每次取当前节点左右儿子中的最小值替换当前值，并一直往下调整，直至找到合适 13 这个元素的位置才算调整完成。

显然，在向下调整的过程中，最多只会进行 log(n) 次比较以及覆盖，所以这个操作的复杂度也是 O(logn)。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;
}

//向下调整，重新确定最后一个元素的位置
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) {
        //计算左儿子下标
        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;
        //最小(最大)值往上提
        queue[k] = c;
        k = child;
    }
    queue[k] = key;
}

//通过比较器来比较
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;
}

另外，与删除元素相关的还有一个 remove 方法，它的作用是删除指定元素。很容易想到，对于一个这种结构的容器来说，删除指定元素免不了需要对整个数组进行遍历，实际上它也是这么做的，所以这个操作的复杂度是很高的，复杂度为O(n)，不过这个操作在实际中是很少用到的，源码如下：

public boolean remove(Object o) {
    int i = indexOf(o);
    if (i == -1)
        return false;
    else {
        removeAt(i);
        return true;
    }
}

//遍历数组，获取元素下标
private int indexOf(Object o) {
    if (o != null) {
        for (int i = 0; i < size; i++)
            if (o.equals(queue[i]))
                return i;
    }
    return -1;
}

//删除指定下标的元素，并对堆结构进行调整
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;
}

---

5、获取元素

获取堆顶元素的代码很简单，其实就是直接取 queue[0] 即可，源码如下：

public E peek() {
    return (size == 0) ? null : (E) queue[0];
}

五、总结

1、PriorityQueue 是一个具有优先级的队列，是通过堆来实现的，它保证队头一定是最大/最小的元素，它底层的存储结构是数组。
2、PriorityQueue 添加元素和删除队头元素的复杂度都是 O(logn)，删除指定元素的复杂度为 O(n)，n 指容器中元素的个数。
3、PriorityQueue 默认的大小是 11。
4、PriorityQueue 在扩容的时候会有一个阈值，当容量小于这个阈值时，它的新容量是旧容量的两倍加2；当容量超过这个阈值时，它是以原容量的 50% 去扩容的。