java集合类型源码解析之PriorityQueue

时间:2022-03-09 14:07:51

本来第二篇想解析一下LinkedList,不过扫了一下源码后,觉得LinkedList的实现比较简单,没有什么意思,于是移步PriorityQueue。

PriorityQueue通过数组实现了一个堆数据结构(相当于一棵完全二叉树),元素的优先级可以通过一个Comparator来计算,如果不指定Comparator,那么元素类型应该实现Comparable接口。最终compare得出的最小元素,放在堆的根部。

成员变量

public class PriorityQueue<E>  extends AbstractQueue<E> {
    transient Object[] queue; // non-private to simplify nested class access
    private int size = 0;
    private final Comparator<? super E> comparator;
    transient int modCount = 0; // non-private to simplify nested class access
}

PriorityQueue的成员变量和ArrayList高度类似:

  • queue 元素存储空间数组,代表一棵完全二叉树,索引位置n的节点的左右子树分别为(2n+1)和(2n+2);
  • size 元素数量
  • modCount 队列修改追踪标记
  • comparator 优先级比较器,可以为null

堆算法

PriorityQueue最重要的部分就是维护堆的几个方法,它基本实现了《算法导论》介绍的堆算法。

1、插入元素siftDown
假定k位置的左右子树都是堆,siftDown方法把元素x插入位置k,然后对这个子堆进行调整,保证以k为根的子树也是堆。视comparator是否存在,siftDown使用siftDownUsingComparator或siftDownComparable,它们只是比较元素的方式不一样,结构是完全一致的,这里就只解读siftDownComparable。

 private void siftDown(int k, E x) {
    if (comparator != null)
        siftDownUsingComparator(k, x);
    else
        siftDownComparable(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;
        queue[k] = c;
        k = child;
    }
    queue[k] = key;
}
  • 先计算一个half位置,因为这个位置之后就是叶节点,不需要继续往下;
  • 计算k的左右子节点,找出compare值最小的节点;
  • 如果就是x,那么k就是插入位置;
  • 如果是子节点,子节点value上移,k指向子节点位置
  • 继续尝试向下探索

2、插入元素siftUp

往k位置插入值x,siftUp采用的方式是往树根方向移动,将祖先节点compare值比较大的往下交换。
siftup的用途要少一点。

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 = parent;
    }
    queue[k] = key;
}
  • 如果k>0,那么k不是树根,就可以继续往上探索;
  • 找到k的父节点parent,比较x和parent的值
  • 如果x>=parent,说明x放在k,以parent为根是一个子堆;
  • 否则将parent放入k位置,k指向parent,继续探索

3、建堆

@SuppressWarnings("unchecked")
private void heapify() {
    for (int i = (size >>> 1) - 1; i >= 0; i--)
        siftDown(i, (E) queue[i]);
}

从中间索引位置开始(往后的都是叶节点),往树根方向遍历,对每个位置调用siftDown。

  • i的初始值,由于只有一层叶子节点,所以满足左右子树都是子堆的条件,siftDown使得,以i为树根的节点也是子堆;
  • 当i变小时,由于k>i的子树已经是堆,那么必然i的左右子树都是堆,所以siftDown使得,以i为树根的节点也是子堆;
  • 最终siftDown(0,queue[0])使得整棵树成为一个堆。

offer和poll

现在可以来看看queue最常见的两个操作:offer和poll。

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

offer先把一个对象放入队列的末尾,前面的空间检查及增长和ArrayList极为类似。
i就是插入位置,如果i==0,queue是空的,插入就完事了;否则通过siftUp来插入,因为i是叶节点,所以可以认为i子树是一个子堆,siftup保证e会往树根方向,找到一个合适的位置,使整棵树保持了堆的特性。

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

poll取出树根元素作为返回值,然后将最后一个元素取出来,通过siftDown插入到树根位置,前面讲过siftDown的性质,这个操作使得整棵树仍然保持堆特性。

remove操作

public boolean remove(Object o) {
    int i = indexOf(o);
    if (i == -1)
        return false;
    else {
        removeAt(i);
        return true;
    }
}
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;
}

这个private的removeAt操作比较有意思,它执行的操作是删除第i个节点(存储空间的第i个,而不是优先级的第i个,这块容易引起人的误解,所以没有类似的public方法)。

  • 如果删除的是最后一个元素,直接置为null即可;
  • 否则把最后一个元素取出来,插入到i位置
  • 插入的过程是先siftDown,如果siftDown的最终位置就是i,那么说明move比i的子树元素都小,此时再尝试一下siftUp;否则siftUp是不需要执行的;
  • 当siftUp执行的结果是末尾元素,被移动到了i之前,那么返回这个元素,其他情况都返回null

这个返回值也是出人意料,不是返回删除的元素,而是在保持堆特性的过程中,如果有尾部元素被移动到i之前的位置,就返回它。这纯粹是为了帮助PriorityQueue的迭代器实现,下一节马上解释。

迭代器

首先要明确一点,PriorityQueue的迭代器并不按优先级顺序来遍历元素,主要就是按存储顺序来遍历,先看迭代器的成员

private final class Itr implements Iterator<E> {
    private int cursor = 0;
    private int lastRet = -1;
    private int expectedModCount = modCount;

    private ArrayDeque<E> forgetMeNot = null;
    private E lastRetElt = null;
}

cursor、lastRet、expectedModCount的作用和ArrayList的迭代器完全一致;但是多出来的forgetMeNot和lastRetElt让人有点莫民奇妙。

再看看remove方法的实现:

public void remove() {
    if (expectedModCount != modCount)
        throw new ConcurrentModificationException();
    if (lastRet != -1) {
        E moved = PriorityQueue.this.removeAt(lastRet);
        lastRet = -1;
        if (moved == null)
            cursor--;
        else {
            if (forgetMeNot == null)
                forgetMeNot = new ArrayDeque<>();
            forgetMeNot.add(moved);
        }
    } else if (lastRetElt != null) {
        PriorityQueue.this.removeEq(lastRetElt);
        lastRetElt = null;
    } else {
        throw new IllegalStateException();
    }
    expectedModCount = modCount;
}

如果lastRet有效,那么调用PriorityQueued.removeAt(lastRet)来删除元素,通过上一节我们知道,removeAt方法可能导致某个元素从末尾被移动到lastRet前面,这样的话,迭代器就会丢失这个元素。为了解决这个问题,迭代器把这个元素放到了一个临时ArrayDeque里面。

这样如果lastRet没有指向有效的元素,那么有可能正在遍历ArrayDeque里面的元素,此时通过lastRetElt来指向。

再看next方法就很容易明白了

public E next() {
    if (expectedModCount != modCount)
        throw new ConcurrentModificationException();
    if (cursor < size)
        return (E) queue[lastRet = cursor++];
    if (forgetMeNot != null) {
        lastRet = -1;
        lastRetElt = forgetMeNot.poll();
        if (lastRetElt != null)
            return lastRetElt;
    }
    throw new NoSuchElementException();
}

先顺着cursor遍历,再把forgetMeNot里面的元素遍历一遍。