com.google.common.collect.TreeMultiset Maven / Gradle / Ivy
/*
* Copyright (C) 2007 The Guava Authors
*
* 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 com.google.common.collect;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.CollectPreconditions.checkNonnegative;
import static com.google.common.collect.NullnessCasts.uncheckedCastNullableTToT;
import static java.util.Objects.requireNonNull;
import com.google.common.annotations.GwtCompatible;
import com.google.common.annotations.GwtIncompatible;
import com.google.common.base.MoreObjects;
import com.google.common.primitives.Ints;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.util.Comparator;
import java.util.ConcurrentModificationException;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.function.ObjIntConsumer;
import javax.annotation.CheckForNull;
import org.checkerframework.checker.nullness.qual.Nullable;
/**
* A multiset which maintains the ordering of its elements, according to either their natural order
* or an explicit {@link Comparator}. In all cases, this implementation uses {@link
* Comparable#compareTo} or {@link Comparator#compare} instead of {@link Object#equals} to determine
* equivalence of instances.
*
* Warning: The comparison must be consistent with equals as explained by the
* {@link Comparable} class specification. Otherwise, the resulting multiset will violate the {@link
* java.util.Collection} contract, which is specified in terms of {@link Object#equals}.
*
*
See the Guava User Guide article on {@code Multiset}.
*
* @author Louis Wasserman
* @author Jared Levy
* @since 2.0
*/
@GwtCompatible(emulated = true)
@ElementTypesAreNonnullByDefault
public final class TreeMultiset extends AbstractSortedMultiset
implements Serializable {
/**
* Creates a new, empty multiset, sorted according to the elements' natural order. All elements
* inserted into the multiset must implement the {@code Comparable} interface. Furthermore, all
* such elements must be mutually comparable: {@code e1.compareTo(e2)} must not throw a
* {@code ClassCastException} for any elements {@code e1} and {@code e2} in the multiset. If the
* user attempts to add an element to the multiset that violates this constraint (for example, the
* user attempts to add a string element to a set whose elements are integers), the {@code
* add(Object)} call will throw a {@code ClassCastException}.
*
* The type specification is {@code }, instead of the more specific
* {@code >}, to support classes defined without generics.
*/
public static TreeMultiset create() {
return new TreeMultiset(Ordering.natural());
}
/**
* Creates a new, empty multiset, sorted according to the specified comparator. All elements
* inserted into the multiset must be mutually comparable by the specified comparator:
* {@code comparator.compare(e1, e2)} must not throw a {@code ClassCastException} for any elements
* {@code e1} and {@code e2} in the multiset. If the user attempts to add an element to the
* multiset that violates this constraint, the {@code add(Object)} call will throw a {@code
* ClassCastException}.
*
* @param comparator the comparator that will be used to sort this multiset. A null value
* indicates that the elements' natural ordering should be used.
*/
@SuppressWarnings("unchecked")
public static TreeMultiset create(
@CheckForNull Comparator comparator) {
return (comparator == null)
? new TreeMultiset((Comparator) Ordering.natural())
: new TreeMultiset(comparator);
}
/**
* Creates an empty multiset containing the given initial elements, sorted according to the
* elements' natural order.
*
* This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
*
*
The type specification is {@code }, instead of the more specific
* {@code >}, to support classes defined without generics.
*/
public static TreeMultiset create(Iterable elements) {
TreeMultiset multiset = create();
Iterables.addAll(multiset, elements);
return multiset;
}
private final transient Reference> rootReference;
private final transient GeneralRange range;
private final transient AvlNode header;
TreeMultiset(Reference> rootReference, GeneralRange range, AvlNode endLink) {
super(range.comparator());
this.rootReference = rootReference;
this.range = range;
this.header = endLink;
}
TreeMultiset(Comparator comparator) {
super(comparator);
this.range = GeneralRange.all(comparator);
this.header = new AvlNode<>();
successor(header, header);
this.rootReference = new Reference<>();
}
/** A function which can be summed across a subtree. */
private enum Aggregate {
SIZE {
@Override
int nodeAggregate(AvlNode node) {
return node.elemCount;
}
@Override
long treeAggregate(@CheckForNull AvlNode root) {
return (root == null) ? 0 : root.totalCount;
}
},
DISTINCT {
@Override
int nodeAggregate(AvlNode node) {
return 1;
}
@Override
long treeAggregate(@CheckForNull AvlNode root) {
return (root == null) ? 0 : root.distinctElements;
}
};
abstract int nodeAggregate(AvlNode node);
abstract long treeAggregate(@CheckForNull AvlNode root);
}
private long aggregateForEntries(Aggregate aggr) {
AvlNode root = rootReference.get();
long total = aggr.treeAggregate(root);
if (range.hasLowerBound()) {
total -= aggregateBelowRange(aggr, root);
}
if (range.hasUpperBound()) {
total -= aggregateAboveRange(aggr, root);
}
return total;
}
private long aggregateBelowRange(Aggregate aggr, @CheckForNull AvlNode node) {
if (node == null) {
return 0;
}
// The cast is safe because we call this method only if hasLowerBound().
int cmp =
comparator()
.compare(uncheckedCastNullableTToT(range.getLowerEndpoint()), node.getElement());
if (cmp < 0) {
return aggregateBelowRange(aggr, node.left);
} else if (cmp == 0) {
switch (range.getLowerBoundType()) {
case OPEN:
return aggr.nodeAggregate(node) + aggr.treeAggregate(node.left);
case CLOSED:
return aggr.treeAggregate(node.left);
default:
throw new AssertionError();
}
} else {
return aggr.treeAggregate(node.left)
+ aggr.nodeAggregate(node)
+ aggregateBelowRange(aggr, node.right);
}
}
private long aggregateAboveRange(Aggregate aggr, @CheckForNull AvlNode node) {
if (node == null) {
return 0;
}
// The cast is safe because we call this method only if hasUpperBound().
int cmp =
comparator()
.compare(uncheckedCastNullableTToT(range.getUpperEndpoint()), node.getElement());
if (cmp > 0) {
return aggregateAboveRange(aggr, node.right);
} else if (cmp == 0) {
switch (range.getUpperBoundType()) {
case OPEN:
return aggr.nodeAggregate(node) + aggr.treeAggregate(node.right);
case CLOSED:
return aggr.treeAggregate(node.right);
default:
throw new AssertionError();
}
} else {
return aggr.treeAggregate(node.right)
+ aggr.nodeAggregate(node)
+ aggregateAboveRange(aggr, node.left);
}
}
@Override
public int size() {
return Ints.saturatedCast(aggregateForEntries(Aggregate.SIZE));
}
@Override
int distinctElements() {
return Ints.saturatedCast(aggregateForEntries(Aggregate.DISTINCT));
}
static int distinctElements(@CheckForNull AvlNode node) {
return (node == null) ? 0 : node.distinctElements;
}
@Override
public int count(@CheckForNull Object element) {
try {
@SuppressWarnings("unchecked")
E e = (E) element;
AvlNode root = rootReference.get();
if (!range.contains(e) || root == null) {
return 0;
}
return root.count(comparator(), e);
} catch (ClassCastException | NullPointerException e) {
return 0;
}
}
@CanIgnoreReturnValue
@Override
public int add(@ParametricNullness E element, int occurrences) {
checkNonnegative(occurrences, "occurrences");
if (occurrences == 0) {
return count(element);
}
checkArgument(range.contains(element));
AvlNode root = rootReference.get();
if (root == null) {
comparator().compare(element, element);
AvlNode newRoot = new AvlNode(element, occurrences);
successor(header, newRoot, header);
rootReference.checkAndSet(root, newRoot);
return 0;
}
int[] result = new int[1]; // used as a mutable int reference to hold result
AvlNode newRoot = root.add(comparator(), element, occurrences, result);
rootReference.checkAndSet(root, newRoot);
return result[0];
}
@CanIgnoreReturnValue
@Override
public int remove(@CheckForNull Object element, int occurrences) {
checkNonnegative(occurrences, "occurrences");
if (occurrences == 0) {
return count(element);
}
AvlNode root = rootReference.get();
int[] result = new int[1]; // used as a mutable int reference to hold result
AvlNode newRoot;
try {
@SuppressWarnings("unchecked")
E e = (E) element;
if (!range.contains(e) || root == null) {
return 0;
}
newRoot = root.remove(comparator(), e, occurrences, result);
} catch (ClassCastException | NullPointerException e) {
return 0;
}
rootReference.checkAndSet(root, newRoot);
return result[0];
}
@CanIgnoreReturnValue
@Override
public int setCount(@ParametricNullness E element, int count) {
checkNonnegative(count, "count");
if (!range.contains(element)) {
checkArgument(count == 0);
return 0;
}
AvlNode root = rootReference.get();
if (root == null) {
if (count > 0) {
add(element, count);
}
return 0;
}
int[] result = new int[1]; // used as a mutable int reference to hold result
AvlNode newRoot = root.setCount(comparator(), element, count, result);
rootReference.checkAndSet(root, newRoot);
return result[0];
}
@CanIgnoreReturnValue
@Override
public boolean setCount(@ParametricNullness E element, int oldCount, int newCount) {
checkNonnegative(newCount, "newCount");
checkNonnegative(oldCount, "oldCount");
checkArgument(range.contains(element));
AvlNode root = rootReference.get();
if (root == null) {
if (oldCount == 0) {
if (newCount > 0) {
add(element, newCount);
}
return true;
} else {
return false;
}
}
int[] result = new int[1]; // used as a mutable int reference to hold result
AvlNode newRoot = root.setCount(comparator(), element, oldCount, newCount, result);
rootReference.checkAndSet(root, newRoot);
return result[0] == oldCount;
}
@Override
public void clear() {
if (!range.hasLowerBound() && !range.hasUpperBound()) {
// We can do this in O(n) rather than removing one by one, which could force rebalancing.
for (AvlNode current = header.succ(); current != header; ) {
AvlNode next = current.succ();
current.elemCount = 0;
// Also clear these fields so that one deleted Entry doesn't retain all elements.
current.left = null;
current.right = null;
current.pred = null;
current.succ = null;
current = next;
}
successor(header, header);
rootReference.clear();
} else {
// TODO(cpovirk): Perhaps we can optimize in this case, too?
Iterators.clear(entryIterator());
}
}
private Entry wrapEntry(final AvlNode baseEntry) {
return new Multisets.AbstractEntry() {
@Override
@ParametricNullness
public E getElement() {
return baseEntry.getElement();
}
@Override
public int getCount() {
int result = baseEntry.getCount();
if (result == 0) {
return count(getElement());
} else {
return result;
}
}
};
}
/** Returns the first node in the tree that is in range. */
@CheckForNull
private AvlNode firstNode() {
AvlNode root = rootReference.get();
if (root == null) {
return null;
}
AvlNode node;
if (range.hasLowerBound()) {
// The cast is safe because of the hasLowerBound check.
E endpoint = uncheckedCastNullableTToT(range.getLowerEndpoint());
node = root.ceiling(comparator(), endpoint);
if (node == null) {
return null;
}
if (range.getLowerBoundType() == BoundType.OPEN
&& comparator().compare(endpoint, node.getElement()) == 0) {
node = node.succ();
}
} else {
node = header.succ();
}
return (node == header || !range.contains(node.getElement())) ? null : node;
}
@CheckForNull
private AvlNode lastNode() {
AvlNode root = rootReference.get();
if (root == null) {
return null;
}
AvlNode node;
if (range.hasUpperBound()) {
// The cast is safe because of the hasUpperBound check.
E endpoint = uncheckedCastNullableTToT(range.getUpperEndpoint());
node = root.floor(comparator(), endpoint);
if (node == null) {
return null;
}
if (range.getUpperBoundType() == BoundType.OPEN
&& comparator().compare(endpoint, node.getElement()) == 0) {
node = node.pred();
}
} else {
node = header.pred();
}
return (node == header || !range.contains(node.getElement())) ? null : node;
}
@Override
Iterator elementIterator() {
return Multisets.elementIterator(entryIterator());
}
@Override
Iterator> entryIterator() {
return new Iterator>() {
@CheckForNull AvlNode current = firstNode();
@CheckForNull Entry prevEntry;
@Override
public boolean hasNext() {
if (current == null) {
return false;
} else if (range.tooHigh(current.getElement())) {
current = null;
return false;
} else {
return true;
}
}
@Override
public Entry next() {
if (!hasNext()) {
throw new NoSuchElementException();
}
// requireNonNull is safe because current is only nulled out after iteration is complete.
Entry result = wrapEntry(requireNonNull(current));
prevEntry = result;
if (current.succ() == header) {
current = null;
} else {
current = current.succ();
}
return result;
}
@Override
public void remove() {
checkState(prevEntry != null, "no calls to next() since the last call to remove()");
setCount(prevEntry.getElement(), 0);
prevEntry = null;
}
};
}
@Override
Iterator> descendingEntryIterator() {
return new Iterator>() {
@CheckForNull AvlNode current = lastNode();
@CheckForNull Entry prevEntry = null;
@Override
public boolean hasNext() {
if (current == null) {
return false;
} else if (range.tooLow(current.getElement())) {
current = null;
return false;
} else {
return true;
}
}
@Override
public Entry next() {
if (!hasNext()) {
throw new NoSuchElementException();
}
// requireNonNull is safe because current is only nulled out after iteration is complete.
requireNonNull(current);
Entry result = wrapEntry(current);
prevEntry = result;
if (current.pred() == header) {
current = null;
} else {
current = current.pred();
}
return result;
}
@Override
public void remove() {
checkState(prevEntry != null, "no calls to next() since the last call to remove()");
setCount(prevEntry.getElement(), 0);
prevEntry = null;
}
};
}
@Override
public void forEachEntry(ObjIntConsumer action) {
checkNotNull(action);
for (AvlNode node = firstNode();
node != header && node != null && !range.tooHigh(node.getElement());
node = node.succ()) {
action.accept(node.getElement(), node.getCount());
}
}
@Override
public Iterator iterator() {
return Multisets.iteratorImpl(this);
}
@Override
public SortedMultiset headMultiset(@ParametricNullness E upperBound, BoundType boundType) {
return new TreeMultiset(
rootReference,
range.intersect(GeneralRange.upTo(comparator(), upperBound, boundType)),
header);
}
@Override
public SortedMultiset tailMultiset(@ParametricNullness E lowerBound, BoundType boundType) {
return new TreeMultiset(
rootReference,
range.intersect(GeneralRange.downTo(comparator(), lowerBound, boundType)),
header);
}
private static final class Reference {
@CheckForNull private T value;
@CheckForNull
public T get() {
return value;
}
public void checkAndSet(@CheckForNull T expected, @CheckForNull T newValue) {
if (value != expected) {
throw new ConcurrentModificationException();
}
value = newValue;
}
void clear() {
value = null;
}
}
private static final class AvlNode {
/*
* For "normal" nodes, the type of this field is `E`, not `@Nullable E` (though note that E is a
* type that can include null, as in a TreeMultiset<@Nullable String>).
*
* For the header node, though, this field contains `null`, regardless of the type of the
* multiset.
*
* Most code that operates on an AvlNode never operates on the header node. Such code can access
* the elem field without a null check by calling getElement().
*/
@CheckForNull private final E elem;
// elemCount is 0 iff this node has been deleted.
private int elemCount;
private int distinctElements;
private long totalCount;
private int height;
@CheckForNull private AvlNode left;
@CheckForNull private AvlNode right;
/*
* pred and succ are nullable after construction, but we always call successor() to initialize
* them immediately thereafter.
*
* They may be subsequently nulled out by TreeMultiset.clear(). I think that the only place that
* we can reference a node whose fields have been cleared is inside the iterator (and presumably
* only under concurrent modification).
*
* To access these fields when you know that they are not null, call the pred() and succ()
* methods, which perform null checks before returning the fields.
*/
@CheckForNull private AvlNode pred;
@CheckForNull private AvlNode succ;
AvlNode(@ParametricNullness E elem, int elemCount) {
checkArgument(elemCount > 0);
this.elem = elem;
this.elemCount = elemCount;
this.totalCount = elemCount;
this.distinctElements = 1;
this.height = 1;
this.left = null;
this.right = null;
}
/** Constructor for the header node. */
AvlNode() {
this.elem = null;
this.elemCount = 1;
}
// For discussion of pred() and succ(), see the comment on the pred and succ fields.
private AvlNode pred() {
return requireNonNull(pred);
}
private AvlNode succ() {
return requireNonNull(succ);
}
int count(Comparator comparator, @ParametricNullness E e) {
int cmp = comparator.compare(e, getElement());
if (cmp < 0) {
return (left == null) ? 0 : left.count(comparator, e);
} else if (cmp > 0) {
return (right == null) ? 0 : right.count(comparator, e);
} else {
return elemCount;
}
}
private AvlNode addRightChild(@ParametricNullness E e, int count) {
right = new AvlNode(e, count);
successor(this, right, succ());
height = Math.max(2, height);
distinctElements++;
totalCount += count;
return this;
}
private AvlNode addLeftChild(@ParametricNullness E e, int count) {
left = new AvlNode(e, count);
successor(pred(), left, this);
height = Math.max(2, height);
distinctElements++;
totalCount += count;
return this;
}
AvlNode add(
Comparator comparator, @ParametricNullness E e, int count, int[] result) {
/*
* It speeds things up considerably to unconditionally add count to totalCount here,
* but that destroys failure atomicity in the case of count overflow. =(
*/
int cmp = comparator.compare(e, getElement());
if (cmp < 0) {
AvlNode initLeft = left;
if (initLeft == null) {
result[0] = 0;
return addLeftChild(e, count);
}
int initHeight = initLeft.height;
left = initLeft.add(comparator, e, count, result);
if (result[0] == 0) {
distinctElements++;
}
this.totalCount += count;
return (left.height == initHeight) ? this : rebalance();
} else if (cmp > 0) {
AvlNode initRight = right;
if (initRight == null) {
result[0] = 0;
return addRightChild(e, count);
}
int initHeight = initRight.height;
right = initRight.add(comparator, e, count, result);
if (result[0] == 0) {
distinctElements++;
}
this.totalCount += count;
return (right.height == initHeight) ? this : rebalance();
}
// adding count to me! No rebalance possible.
result[0] = elemCount;
long resultCount = (long) elemCount + count;
checkArgument(resultCount <= Integer.MAX_VALUE);
this.elemCount += count;
this.totalCount += count;
return this;
}
@CheckForNull
AvlNode remove(
Comparator comparator, @ParametricNullness E e, int count, int[] result) {
int cmp = comparator.compare(e, getElement());
if (cmp < 0) {
AvlNode initLeft = left;
if (initLeft == null) {
result[0] = 0;
return this;
}
left = initLeft.remove(comparator, e, count, result);
if (result[0] > 0) {
if (count >= result[0]) {
this.distinctElements--;
this.totalCount -= result[0];
} else {
this.totalCount -= count;
}
}
return (result[0] == 0) ? this : rebalance();
} else if (cmp > 0) {
AvlNode initRight = right;
if (initRight == null) {
result[0] = 0;
return this;
}
right = initRight.remove(comparator, e, count, result);
if (result[0] > 0) {
if (count >= result[0]) {
this.distinctElements--;
this.totalCount -= result[0];
} else {
this.totalCount -= count;
}
}
return rebalance();
}
// removing count from me!
result[0] = elemCount;
if (count >= elemCount) {
return deleteMe();
} else {
this.elemCount -= count;
this.totalCount -= count;
return this;
}
}
@CheckForNull
AvlNode setCount(
Comparator comparator, @ParametricNullness E e, int count, int[] result) {
int cmp = comparator.compare(e, getElement());
if (cmp < 0) {
AvlNode initLeft = left;
if (initLeft == null) {
result[0] = 0;
return (count > 0) ? addLeftChild(e, count) : this;
}
left = initLeft.setCount(comparator, e, count, result);
if (count == 0 && result[0] != 0) {
this.distinctElements--;
} else if (count > 0 && result[0] == 0) {
this.distinctElements++;
}
this.totalCount += count - result[0];
return rebalance();
} else if (cmp > 0) {
AvlNode initRight = right;
if (initRight == null) {
result[0] = 0;
return (count > 0) ? addRightChild(e, count) : this;
}
right = initRight.setCount(comparator, e, count, result);
if (count == 0 && result[0] != 0) {
this.distinctElements--;
} else if (count > 0 && result[0] == 0) {
this.distinctElements++;
}
this.totalCount += count - result[0];
return rebalance();
}
// setting my count
result[0] = elemCount;
if (count == 0) {
return deleteMe();
}
this.totalCount += count - elemCount;
this.elemCount = count;
return this;
}
@CheckForNull
AvlNode setCount(
Comparator comparator,
@ParametricNullness E e,
int expectedCount,
int newCount,
int[] result) {
int cmp = comparator.compare(e, getElement());
if (cmp < 0) {
AvlNode initLeft = left;
if (initLeft == null) {
result[0] = 0;
if (expectedCount == 0 && newCount > 0) {
return addLeftChild(e, newCount);
}
return this;
}
left = initLeft.setCount(comparator, e, expectedCount, newCount, result);
if (result[0] == expectedCount) {
if (newCount == 0 && result[0] != 0) {
this.distinctElements--;
} else if (newCount > 0 && result[0] == 0) {
this.distinctElements++;
}
this.totalCount += newCount - result[0];
}
return rebalance();
} else if (cmp > 0) {
AvlNode initRight = right;
if (initRight == null) {
result[0] = 0;
if (expectedCount == 0 && newCount > 0) {
return addRightChild(e, newCount);
}
return this;
}
right = initRight.setCount(comparator, e, expectedCount, newCount, result);
if (result[0] == expectedCount) {
if (newCount == 0 && result[0] != 0) {
this.distinctElements--;
} else if (newCount > 0 && result[0] == 0) {
this.distinctElements++;
}
this.totalCount += newCount - result[0];
}
return rebalance();
}
// setting my count
result[0] = elemCount;
if (expectedCount == elemCount) {
if (newCount == 0) {
return deleteMe();
}
this.totalCount += newCount - elemCount;
this.elemCount = newCount;
}
return this;
}
@CheckForNull
private AvlNode deleteMe() {
int oldElemCount = this.elemCount;
this.elemCount = 0;
successor(pred(), succ());
if (left == null) {
return right;
} else if (right == null) {
return left;
} else if (left.height >= right.height) {
AvlNode newTop = pred();
// newTop is the maximum node in my left subtree
newTop.left = left.removeMax(newTop);
newTop.right = right;
newTop.distinctElements = distinctElements - 1;
newTop.totalCount = totalCount - oldElemCount;
return newTop.rebalance();
} else {
AvlNode newTop = succ();
newTop.right = right.removeMin(newTop);
newTop.left = left;
newTop.distinctElements = distinctElements - 1;
newTop.totalCount = totalCount - oldElemCount;
return newTop.rebalance();
}
}
// Removes the minimum node from this subtree to be reused elsewhere
@CheckForNull
private AvlNode removeMin(AvlNode node) {
if (left == null) {
return right;
} else {
left = left.removeMin(node);
distinctElements--;
totalCount -= node.elemCount;
return rebalance();
}
}
// Removes the maximum node from this subtree to be reused elsewhere
@CheckForNull
private AvlNode removeMax(AvlNode node) {
if (right == null) {
return left;
} else {
right = right.removeMax(node);
distinctElements--;
totalCount -= node.elemCount;
return rebalance();
}
}
private void recomputeMultiset() {
this.distinctElements =
1 + TreeMultiset.distinctElements(left) + TreeMultiset.distinctElements(right);
this.totalCount = elemCount + totalCount(left) + totalCount(right);
}
private void recomputeHeight() {
this.height = 1 + Math.max(height(left), height(right));
}
private void recompute() {
recomputeMultiset();
recomputeHeight();
}
private AvlNode rebalance() {
switch (balanceFactor()) {
case -2:
// requireNonNull is safe because right must exist in order to get a negative factor.
requireNonNull(right);
if (right.balanceFactor() > 0) {
right = right.rotateRight();
}
return rotateLeft();
case 2:
// requireNonNull is safe because left must exist in order to get a positive factor.
requireNonNull(left);
if (left.balanceFactor() < 0) {
left = left.rotateLeft();
}
return rotateRight();
default:
recomputeHeight();
return this;
}
}
private int balanceFactor() {
return height(left) - height(right);
}
private AvlNode rotateLeft() {
checkState(right != null);
AvlNode newTop = right;
this.right = newTop.left;
newTop.left = this;
newTop.totalCount = this.totalCount;
newTop.distinctElements = this.distinctElements;
this.recompute();
newTop.recomputeHeight();
return newTop;
}
private AvlNode rotateRight() {
checkState(left != null);
AvlNode newTop = left;
this.left = newTop.right;
newTop.right = this;
newTop.totalCount = this.totalCount;
newTop.distinctElements = this.distinctElements;
this.recompute();
newTop.recomputeHeight();
return newTop;
}
private static long totalCount(@CheckForNull AvlNode node) {
return (node == null) ? 0 : node.totalCount;
}
private static int height(@CheckForNull AvlNode node) {
return (node == null) ? 0 : node.height;
}
@CheckForNull
private AvlNode ceiling(Comparator comparator, @ParametricNullness E e) {
int cmp = comparator.compare(e, getElement());
if (cmp < 0) {
return (left == null) ? this : MoreObjects.firstNonNull(left.ceiling(comparator, e), this);
} else if (cmp == 0) {
return this;
} else {
return (right == null) ? null : right.ceiling(comparator, e);
}
}
@CheckForNull
private AvlNode floor(Comparator comparator, @ParametricNullness E e) {
int cmp = comparator.compare(e, getElement());
if (cmp > 0) {
return (right == null) ? this : MoreObjects.firstNonNull(right.floor(comparator, e), this);
} else if (cmp == 0) {
return this;
} else {
return (left == null) ? null : left.floor(comparator, e);
}
}
@ParametricNullness
E getElement() {
// For discussion of this cast, see the comment on the elem field.
return uncheckedCastNullableTToT(elem);
}
int getCount() {
return elemCount;
}
@Override
public String toString() {
return Multisets.immutableEntry(getElement(), getCount()).toString();
}
}
private static void successor(AvlNode a, AvlNode b) {
a.succ = b;
b.pred = a;
}
private static void successor(
AvlNode a, AvlNode b, AvlNode c) {
successor(a, b);
successor(b, c);
}
/*
* TODO(jlevy): Decide whether entrySet() should return entries with an equals() method that
* calls the comparator to compare the two keys. If that change is made,
* AbstractMultiset.equals() can simply check whether two multisets have equal entry sets.
*/
/**
* @serialData the comparator, the number of distinct elements, the first element, its count, the
* second element, its count, and so on
*/
@GwtIncompatible // java.io.ObjectOutputStream
private void writeObject(ObjectOutputStream stream) throws IOException {
stream.defaultWriteObject();
stream.writeObject(elementSet().comparator());
Serialization.writeMultiset(this, stream);
}
@GwtIncompatible // java.io.ObjectInputStream
private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
stream.defaultReadObject();
@SuppressWarnings("unchecked")
// reading data stored by writeObject
Comparator comparator = (Comparator) stream.readObject();
Serialization.getFieldSetter(AbstractSortedMultiset.class, "comparator").set(this, comparator);
Serialization.getFieldSetter(TreeMultiset.class, "range")
.set(this, GeneralRange.all(comparator));
Serialization.getFieldSetter(TreeMultiset.class, "rootReference")
.set(this, new Reference>());
AvlNode header = new AvlNode<>();
Serialization.getFieldSetter(TreeMultiset.class, "header").set(this, header);
successor(header, header);
Serialization.populateMultiset(this, stream);
}
@GwtIncompatible // not needed in emulated source
private static final long serialVersionUID = 1;
}