@CanIgnoreReturnValue
@Override
public int remove(@Nullable Object element, int occurrences) {
checkNonnegative(occurrences, "occurrences");
if (occurrences == 0) {
return count(element);
}
AvlNode<E> root = rootReference.get();
int[] result = new int[1]; // used as a mutable int reference to hold result
AvlNode<E> 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 e) {
return 0;
} catch (NullPointerException e) {
return 0;
}
rootReference.checkAndSet(root, newRoot);
return result[0];
}
/**
* Rotate binary tree node with right child. For AVL trees, this is a single rotation for case 4.
* Update heights, then return new root.
*/
private static AvlNode rotateWithRightChild(AvlNode k1) {
AvlNode k2 = k1.right;
k1.right = k2.left;
k2.left = k1;
k1.height = max(height(k1.left), height(k1.right)) + 1;
k2.height = max(height(k2.right), k1.height) + 1;
return k2;
}
private AvlNode<T> rotateWithRightChild(AvlNode<T> k2) {
AvlNode<T> k1 = k2.right;
k2.right = k1.left;
k1.left = k2;
k2.height = Math.max(height(k2.left), height(k2.right)) + 1;
k1.height = Math.max(k2.height, height(k1.right)) + 1;
return k1;
}
public int count(Comparator<? super E> comparator, E e) {
int cmp = comparator.compare(e, elem);
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<E> rotateRight() {
checkState(left != null);
AvlNode<E> newTop = left;
this.left = newTop.right;
newTop.right = this;
newTop.totalCount = this.totalCount;
newTop.distinctElements = this.distinctElements;
this.recompute();
newTop.recomputeHeight();
return newTop;
}
@Nullable
private AvlNode<E> floor(Comparator<? super E> comparator, E e) {
int cmp = comparator.compare(e, elem);
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);
}
}
AvlNode<E> remove(Comparator<? super E> comparator, @Nullable E e, int count, int[] result) {
int cmp = comparator.compare(e, elem);
if (cmp < 0) {
AvlNode<E> 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<E> 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;
}
}
@Override
public int count(@Nullable Object element) {
try {
@SuppressWarnings("unchecked")
E e = (E) element;
AvlNode<E> root = rootReference.get();
if (!range.contains(e) || root == null) {
return 0;
}
return root.count(comparator(), e);
} catch (ClassCastException e) {
return 0;
} catch (NullPointerException e) {
return 0;
}
}
/**
* Internal method to insert into a subtree.
*
* @param x the item to insert.
* @param t the node that roots the tree.
* @return the new root.
*/
private AvlNode insert(Comparable x, AvlNode t) {
if (t == null) t = new AvlNode(x, null, null);
else if (x.compareTo(t.element) < 0) {
t.left = insert(x, t.left);
if (height(t.left) - height(t.right) == 2)
if (x.compareTo(t.left.element) < 0) t = rotateWithLeftChild(t);
else t = doubleWithLeftChild(t);
} else if (x.compareTo(t.element) > 0) {
t.right = insert(x, t.right);
if (height(t.right) - height(t.left) == 2)
if (x.compareTo(t.right.element) > 0) t = rotateWithRightChild(t);
else t = doubleWithRightChild(t);
} else ; // Duplicate; do nothing
t.height = max(height(t.left), height(t.right)) + 1;
return t;
}
AvlNode<E> setCount(Comparator<? super E> comparator, @Nullable E e, int count, int[] result) {
int cmp = comparator.compare(e, elem);
if (cmp < 0) {
AvlNode<E> 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<E> 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;
}
public AvlNode<T> insert(T x, AvlNode<T> t) {
if (t == null) return new AvlNode<>(x, null, null);
int compareResult = compare(x, t.element);
if (compareResult < 0) {
t.left = insert(x, t.left);
if (height(t.left) - height(t.right) == 2)
if (compare(x, t.left.element) < 0) t = rotateWithLeftChild(t);
else t = doubleWithLeftChild(t);
} else if (compareResult > 0) {
t.right = insert(x, t.right);
if (height(t.right) - height(t.left) == 2)
if (compare(x, t.right.element) > 0) t = rotateWithRightChild(t);
else t = doubleWithRightChild(t);
} else ;
t.height = Math.max(height(t.left), height(t.right)) + 1;
return t;
}
private AvlNode<E> rebalance() {
switch (balanceFactor()) {
case -2:
if (right.balanceFactor() > 0) {
right = right.rotateRight();
}
return rotateLeft();
case 2:
if (left.balanceFactor() < 0) {
left = left.rotateLeft();
}
return rotateRight();
default:
recomputeHeight();
return this;
}
}
AvlNode<E> add(Comparator<? super E> comparator, @Nullable 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, elem);
if (cmp < 0) {
AvlNode<E> 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<E> 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;
}
// Removes the minimum node from this subtree to be reused elsewhere
private AvlNode<E> removeMin(AvlNode<E> 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
private AvlNode<E> removeMax(AvlNode<E> node) {
if (right == null) {
return left;
} else {
right = right.removeMax(node);
distinctElements--;
totalCount -= node.elemCount;
return rebalance();
}
}
@CanIgnoreReturnValue
@Override
public int add(@Nullable E element, int occurrences) {
checkNonnegative(occurrences, "occurrences");
if (occurrences == 0) {
return count(element);
}
checkArgument(range.contains(element));
AvlNode<E> root = rootReference.get();
if (root == null) {
comparator().compare(element, element);
AvlNode<E> newRoot = new AvlNode<E>(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<E> newRoot = root.add(comparator(), element, occurrences, result);
rootReference.checkAndSet(root, newRoot);
return result[0];
}
@CanIgnoreReturnValue
@Override
public int setCount(@Nullable E element, int count) {
checkNonnegative(count, "count");
if (!range.contains(element)) {
checkArgument(count == 0);
return 0;
}
AvlNode<E> 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<E> newRoot = root.setCount(comparator(), element, count, result);
rootReference.checkAndSet(root, newRoot);
return result[0];
}
@Nullable
private AvlNode<E> lastNode() {
AvlNode<E> root = rootReference.get();
if (root == null) {
return null;
}
AvlNode<E> node;
if (range.hasUpperBound()) {
E endpoint = range.getUpperEndpoint();
node = rootReference.get().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;
}
@CanIgnoreReturnValue
@Override
public boolean setCount(@Nullable E element, int oldCount, int newCount) {
checkNonnegative(newCount, "newCount");
checkNonnegative(oldCount, "oldCount");
checkArgument(range.contains(element));
AvlNode<E> 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<E> newRoot = root.setCount(comparator(), element, oldCount, newCount, result);
rootReference.checkAndSet(root, newRoot);
return result[0] == oldCount;
}
private static <T> void successor(AvlNode<T> a, AvlNode<T> b) {
a.succ = b;
b.pred = a;
}
/**
* Double rotate binary tree node: first left child with its right child; then node k3 with new
* left child. For AVL trees, this is a double rotation for case 2. Update heights, then return
* new root.
*/
private static AvlNode doubleWithLeftChild(AvlNode k3) {
k3.left = rotateWithRightChild(k3.left);
return rotateWithLeftChild(k3);
}
/**
* Double rotate binary tree node: first right child with its left child; then node k1 with new
* right child. For AVL trees, this is a double rotation for case 3. Update heights, then return
* new root.
*/
private static AvlNode doubleWithRightChild(AvlNode k1) {
k1.right = rotateWithLeftChild(k1.right);
return rotateWithRightChild(k1);
}
private AvlNode<T> doubleWithRightChild(AvlNode<T> k3) {
k3.right = rotateWithLeftChild(k3.right);
return rotateWithRightChild(k3);
}
private AvlNode<E> 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<E> 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<E> newTop = succ;
newTop.right = right.removeMin(newTop);
newTop.left = left;
newTop.distinctElements = distinctElements - 1;
newTop.totalCount = totalCount - oldElemCount;
return newTop.rebalance();
}
}
