blob: a493491c01f4fb7a62bfcb9169919721a8a17179 [file] [log] [blame]
/*
* Copyright 2008 Google Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not
* use this file except in compliance with the License. You may obtain a copy of
* the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations under
* the License.
*/
package java.util;
import static javaemul.internal.InternalPreconditions.checkArgument;
import static javaemul.internal.InternalPreconditions.checkCriticalNotNull;
import static javaemul.internal.InternalPreconditions.checkElement;
import static javaemul.internal.InternalPreconditions.checkNotNull;
import static javaemul.internal.InternalPreconditions.checkState;
/**
* An unbounded priority queue based on a priority heap.
* See <a href="https://docs.oracle.com/javase/8/docs/api/java/util/PriorityQueue.html">
* the official Java API doc</a> for details.
* A priority queue does not permit {@code null} elements.
*
* @param <E> element type.
*/
public class PriorityQueue<E> extends AbstractQueue<E> {
private static final int DEFAULT_INITIAL_CAPACITY = 11;
private static int getLeftChild(int node) {
return 2 * node + 1;
}
private static int getParent(int node) {
return (node - 1) / 2;
}
private static int getRightChild(int node) {
return 2 * node + 2;
}
private static boolean isLeaf(int node, int size) {
return node * 2 + 1 >= size;
}
private Comparator<? super E> cmp;
/**
* A heap held in an array. heap[0] is the root of the heap (the smallest
* element), the subtrees of node i are 2*i+1 (left) and 2*i+2 (right). Node i
* is a leaf node if 2*i>=n. Node i's parent, if i>0, is floor((i-1)/2).
*/
private ArrayList<E> heap;
public PriorityQueue() {
this(DEFAULT_INITIAL_CAPACITY);
}
public PriorityQueue(Collection<? extends E> c) {
this(c.size());
addAll(c);
}
public PriorityQueue(int initialCapacity) {
this(initialCapacity, null);
}
public PriorityQueue(int initialCapacity, Comparator<? super E> cmp) {
heap = new ArrayList<>(initialCapacity);
this.cmp = Comparators.nullToNaturalOrder(cmp);
}
public PriorityQueue(Comparator<? super E> comparator) {
this(DEFAULT_INITIAL_CAPACITY, comparator);
}
@SuppressWarnings("unchecked")
public PriorityQueue(PriorityQueue<? extends E> c) {
this(c.size(), (Comparator<? super E>) c.comparator());
addAll(c);
}
@SuppressWarnings("unchecked")
public PriorityQueue(SortedSet<? extends E> c) {
this(c.size(), (Comparator<? super E>) c.comparator());
addAll(c);
}
@Override
public boolean addAll(Collection<? extends E> c) {
checkNotNull(c);
checkArgument(c != this);
int oldSize = heap.size();
for (E e : c) {
heap.add(checkCriticalNotNull(e));
}
if (oldSize != heap.size()) {
makeHeap(0);
return true;
}
return false;
}
@Override
public void clear() {
heap.clear();
}
public Comparator<? super E> comparator() {
return Comparators.naturalOrderToNull(cmp);
}
@Override
public boolean contains(Object o) {
return indexOf(o) != -1;
}
@Override
public Iterator<E> iterator() {
return new Iterator<E>() {
int i = 0, last = -1;
@Override
public boolean hasNext() {
return i < heap.size();
}
@Override
public E next() {
checkElement(hasNext());
last = i++;
return heap.get(last);
}
@Override
public void remove() {
checkState(last != -1);
removeAtIndex(i = last);
last = -1;
}
};
}
@Override
public boolean offer(E e) {
checkCriticalNotNull(e);
int node = heap.size();
heap.add(e);
while (node > 0) {
int childNode = node;
node = getParent(node);
if (cmp.compare(heap.get(node), e) <= 0) {
// parent is smaller, so we have a valid heap
heap.set(childNode, e);
return true;
}
// exchange with parent and try again
heap.set(childNode, heap.get(node));
}
heap.set(node, e);
return true;
}
@Override
public E peek() {
return heap.isEmpty() ? null : heap.get(0);
}
@Override
public E poll() {
E value = peek();
if (value != null) {
removeAtIndex(0);
}
return value;
}
@Override
public boolean remove(Object o) {
int index = indexOf(o);
if (index < 0) {
return false;
}
removeAtIndex(index);
return true;
}
@Override
public boolean removeAll(Collection<?> c) {
if (heap.removeAll(c)) {
makeHeap(0);
return true;
}
return false;
}
@Override
public boolean retainAll(Collection<?> c) {
if (heap.retainAll(c)) {
makeHeap(0);
return true;
}
return false;
}
@Override
public int size() {
return heap.size();
}
@Override
public Spliterator<E> spliterator() {
return Spliterators.spliterator(this, Spliterator.NONNULL);
}
@Override
public Object[] toArray() {
return heap.toArray();
}
@Override
public <T> T[] toArray(T[] a) {
return heap.toArray(a);
}
/**
* Make the subtree rooted at <code>node</code> a valid heap. O(n) time
*
* @param node
*/
private void makeHeap(int node) {
if (isLeaf(node)) {
// leaf node are automatically valid heaps
return;
}
makeHeap(getLeftChild(node)); // make left subtree a heap
// an interior node might not have a right child
int rightChild = getRightChild(node);
if (rightChild < heap.size()) {
makeHeap(rightChild); // make right subtree a heap
}
mergeHeaps(node);
}
/**
* Merge two subheaps into a single heap. O(log n) time
*
* PRECONDITION: both children of <code>node</code> are heaps
*
* @param node the parent of the two subtrees to merge
*/
private void mergeHeaps(int node) {
int heapSize = heap.size();
E value = heap.get(node);
while (!isLeaf(node, heapSize)) {
int smallestChild = getSmallestChild(node, heapSize);
if (cmp.compare(value, heap.get(smallestChild)) < 0) {
// Current node is smaller than the smallest child, so we are done.
break;
}
// Move the smallest child up and iterate using its old slot.
heap.set(node, heap.get(smallestChild));
node = smallestChild;
}
heap.set(node, value);
}
private int getSmallestChild(int node, int heapSize) {
int smallestChild;
int leftChild = getLeftChild(node); // start with left child
int rightChild = leftChild + 1;
smallestChild = leftChild;
if ((rightChild < heapSize)
&& (cmp.compare(heap.get(rightChild), heap.get(leftChild)) < 0)) {
// right child is smaller, go down that path
smallestChild = rightChild;
}
return smallestChild;
}
private int indexOf(Object o) {
return o == null ? -1 : heap.indexOf(o);
}
private boolean isLeaf(int node) {
return isLeaf(node, heap.size());
}
/**
* This method leaves the elements at up to i-1, inclusive, untouched.
* This information is used by PriorityQueue iterator implementation.
*/
private void removeAtIndex(int index) {
// Remove the last element; put it in place of the really removed element.
E lastValue = heap.remove(heap.size() - 1);
// Unless the last element was actually the one we wanted.
if (index < heap.size()) {
// Move last element to the now-empty slot and reheap.
heap.set(index, lastValue);
mergeHeaps(index);
}
}
}