001/* 002 * Copyright (C) 2007 The Guava Authors 003 * 004 * Licensed under the Apache License, Version 2.0 (the "License"); 005 * you may not use this file except in compliance with the License. 006 * You may obtain a copy of the License at 007 * 008 * http://www.apache.org/licenses/LICENSE-2.0 009 * 010 * Unless required by applicable law or agreed to in writing, software 011 * distributed under the License is distributed on an "AS IS" BASIS, 012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 013 * See the License for the specific language governing permissions and 014 * limitations under the License. 015 */ 016 017package com.google.common.collect; 018 019import static com.google.common.base.Preconditions.checkArgument; 020import static com.google.common.base.Preconditions.checkState; 021import static com.google.common.collect.CollectPreconditions.checkNonnegative; 022import static com.google.common.collect.NullnessCasts.uncheckedCastNullableTToT; 023import static java.lang.Math.max; 024import static java.util.Objects.requireNonNull; 025 026import com.google.common.annotations.GwtCompatible; 027import com.google.common.annotations.GwtIncompatible; 028import com.google.common.annotations.J2ktIncompatible; 029import com.google.common.base.MoreObjects; 030import com.google.common.primitives.Ints; 031import com.google.errorprone.annotations.CanIgnoreReturnValue; 032import java.io.IOException; 033import java.io.ObjectInputStream; 034import java.io.ObjectOutputStream; 035import java.io.Serializable; 036import java.util.Comparator; 037import java.util.ConcurrentModificationException; 038import java.util.Iterator; 039import java.util.NoSuchElementException; 040import javax.annotation.CheckForNull; 041import org.checkerframework.checker.nullness.qual.Nullable; 042 043/** 044 * A multiset which maintains the ordering of its elements, according to either their natural order 045 * or an explicit {@link Comparator}. In all cases, this implementation uses {@link 046 * Comparable#compareTo} or {@link Comparator#compare} instead of {@link Object#equals} to determine 047 * equivalence of instances. 048 * 049 * <p><b>Warning:</b> The comparison must be <i>consistent with equals</i> as explained by the 050 * {@link Comparable} class specification. Otherwise, the resulting multiset will violate the {@link 051 * java.util.Collection} contract, which is specified in terms of {@link Object#equals}. 052 * 053 * <p>See the Guava User Guide article on <a href= 054 * "https://github.com/google/guava/wiki/NewCollectionTypesExplained#multiset">{@code Multiset}</a>. 055 * 056 * @author Louis Wasserman 057 * @author Jared Levy 058 * @since 2.0 059 */ 060@GwtCompatible(emulated = true) 061@ElementTypesAreNonnullByDefault 062public final class TreeMultiset<E extends @Nullable Object> extends AbstractSortedMultiset<E> 063 implements Serializable { 064 065 /** 066 * Creates a new, empty multiset, sorted according to the elements' natural order. All elements 067 * inserted into the multiset must implement the {@code Comparable} interface. Furthermore, all 068 * such elements must be <i>mutually comparable</i>: {@code e1.compareTo(e2)} must not throw a 069 * {@code ClassCastException} for any elements {@code e1} and {@code e2} in the multiset. If the 070 * user attempts to add an element to the multiset that violates this constraint (for example, the 071 * user attempts to add a string element to a set whose elements are integers), the {@code 072 * add(Object)} call will throw a {@code ClassCastException}. 073 * 074 * <p>The type specification is {@code <E extends Comparable>}, instead of the more specific 075 * {@code <E extends Comparable<? super E>>}, to support classes defined without generics. 076 */ 077 @SuppressWarnings("rawtypes") // https://github.com/google/guava/issues/989 078 public static <E extends Comparable> TreeMultiset<E> create() { 079 return new TreeMultiset<>(Ordering.natural()); 080 } 081 082 /** 083 * Creates a new, empty multiset, sorted according to the specified comparator. All elements 084 * inserted into the multiset must be <i>mutually comparable</i> by the specified comparator: 085 * {@code comparator.compare(e1, e2)} must not throw a {@code ClassCastException} for any elements 086 * {@code e1} and {@code e2} in the multiset. If the user attempts to add an element to the 087 * multiset that violates this constraint, the {@code add(Object)} call will throw a {@code 088 * ClassCastException}. 089 * 090 * @param comparator the comparator that will be used to sort this multiset. A null value 091 * indicates that the elements' <i>natural ordering</i> should be used. 092 */ 093 @SuppressWarnings("unchecked") 094 public static <E extends @Nullable Object> TreeMultiset<E> create( 095 @CheckForNull Comparator<? super E> comparator) { 096 return (comparator == null) 097 ? new TreeMultiset<E>((Comparator) Ordering.natural()) 098 : new TreeMultiset<E>(comparator); 099 } 100 101 /** 102 * Creates an empty multiset containing the given initial elements, sorted according to the 103 * elements' natural order. 104 * 105 * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}. 106 * 107 * <p>The type specification is {@code <E extends Comparable>}, instead of the more specific 108 * {@code <E extends Comparable<? super E>>}, to support classes defined without generics. 109 */ 110 @SuppressWarnings("rawtypes") // https://github.com/google/guava/issues/989 111 public static <E extends Comparable> TreeMultiset<E> create(Iterable<? extends E> elements) { 112 TreeMultiset<E> multiset = create(); 113 Iterables.addAll(multiset, elements); 114 return multiset; 115 } 116 117 private final transient Reference<AvlNode<E>> rootReference; 118 private final transient GeneralRange<E> range; 119 private final transient AvlNode<E> header; 120 121 TreeMultiset(Reference<AvlNode<E>> rootReference, GeneralRange<E> range, AvlNode<E> endLink) { 122 super(range.comparator()); 123 this.rootReference = rootReference; 124 this.range = range; 125 this.header = endLink; 126 } 127 128 TreeMultiset(Comparator<? super E> comparator) { 129 super(comparator); 130 this.range = GeneralRange.all(comparator); 131 this.header = new AvlNode<>(); 132 successor(header, header); 133 this.rootReference = new Reference<>(); 134 } 135 136 /** A function which can be summed across a subtree. */ 137 private enum Aggregate { 138 SIZE { 139 @Override 140 int nodeAggregate(AvlNode<?> node) { 141 return node.elemCount; 142 } 143 144 @Override 145 long treeAggregate(@CheckForNull AvlNode<?> root) { 146 return (root == null) ? 0 : root.totalCount; 147 } 148 }, 149 DISTINCT { 150 @Override 151 int nodeAggregate(AvlNode<?> node) { 152 return 1; 153 } 154 155 @Override 156 long treeAggregate(@CheckForNull AvlNode<?> root) { 157 return (root == null) ? 0 : root.distinctElements; 158 } 159 }; 160 161 abstract int nodeAggregate(AvlNode<?> node); 162 163 abstract long treeAggregate(@CheckForNull AvlNode<?> root); 164 } 165 166 private long aggregateForEntries(Aggregate aggr) { 167 AvlNode<E> root = rootReference.get(); 168 long total = aggr.treeAggregate(root); 169 if (range.hasLowerBound()) { 170 total -= aggregateBelowRange(aggr, root); 171 } 172 if (range.hasUpperBound()) { 173 total -= aggregateAboveRange(aggr, root); 174 } 175 return total; 176 } 177 178 private long aggregateBelowRange(Aggregate aggr, @CheckForNull AvlNode<E> node) { 179 if (node == null) { 180 return 0; 181 } 182 // The cast is safe because we call this method only if hasLowerBound(). 183 int cmp = 184 comparator() 185 .compare(uncheckedCastNullableTToT(range.getLowerEndpoint()), node.getElement()); 186 if (cmp < 0) { 187 return aggregateBelowRange(aggr, node.left); 188 } else if (cmp == 0) { 189 switch (range.getLowerBoundType()) { 190 case OPEN: 191 return aggr.nodeAggregate(node) + aggr.treeAggregate(node.left); 192 case CLOSED: 193 return aggr.treeAggregate(node.left); 194 } 195 throw new AssertionError(); 196 } else { 197 return aggr.treeAggregate(node.left) 198 + aggr.nodeAggregate(node) 199 + aggregateBelowRange(aggr, node.right); 200 } 201 } 202 203 private long aggregateAboveRange(Aggregate aggr, @CheckForNull AvlNode<E> node) { 204 if (node == null) { 205 return 0; 206 } 207 // The cast is safe because we call this method only if hasUpperBound(). 208 int cmp = 209 comparator() 210 .compare(uncheckedCastNullableTToT(range.getUpperEndpoint()), node.getElement()); 211 if (cmp > 0) { 212 return aggregateAboveRange(aggr, node.right); 213 } else if (cmp == 0) { 214 switch (range.getUpperBoundType()) { 215 case OPEN: 216 return aggr.nodeAggregate(node) + aggr.treeAggregate(node.right); 217 case CLOSED: 218 return aggr.treeAggregate(node.right); 219 } 220 throw new AssertionError(); 221 } else { 222 return aggr.treeAggregate(node.right) 223 + aggr.nodeAggregate(node) 224 + aggregateAboveRange(aggr, node.left); 225 } 226 } 227 228 @Override 229 public int size() { 230 return Ints.saturatedCast(aggregateForEntries(Aggregate.SIZE)); 231 } 232 233 @Override 234 int distinctElements() { 235 return Ints.saturatedCast(aggregateForEntries(Aggregate.DISTINCT)); 236 } 237 238 static int distinctElements(@CheckForNull AvlNode<?> node) { 239 return (node == null) ? 0 : node.distinctElements; 240 } 241 242 @Override 243 public int count(@CheckForNull Object element) { 244 try { 245 @SuppressWarnings("unchecked") 246 E e = (E) element; 247 AvlNode<E> root = rootReference.get(); 248 if (!range.contains(e) || root == null) { 249 return 0; 250 } 251 return root.count(comparator(), e); 252 } catch (ClassCastException | NullPointerException e) { 253 return 0; 254 } 255 } 256 257 @CanIgnoreReturnValue 258 @Override 259 public int add(@ParametricNullness E element, int occurrences) { 260 checkNonnegative(occurrences, "occurrences"); 261 if (occurrences == 0) { 262 return count(element); 263 } 264 checkArgument(range.contains(element)); 265 AvlNode<E> root = rootReference.get(); 266 if (root == null) { 267 int unused = comparator().compare(element, element); 268 AvlNode<E> newRoot = new AvlNode<>(element, occurrences); 269 successor(header, newRoot, header); 270 rootReference.checkAndSet(root, newRoot); 271 return 0; 272 } 273 int[] result = new int[1]; // used as a mutable int reference to hold result 274 AvlNode<E> newRoot = root.add(comparator(), element, occurrences, result); 275 rootReference.checkAndSet(root, newRoot); 276 return result[0]; 277 } 278 279 @CanIgnoreReturnValue 280 @Override 281 public int remove(@CheckForNull Object element, int occurrences) { 282 checkNonnegative(occurrences, "occurrences"); 283 if (occurrences == 0) { 284 return count(element); 285 } 286 AvlNode<E> root = rootReference.get(); 287 int[] result = new int[1]; // used as a mutable int reference to hold result 288 AvlNode<E> newRoot; 289 try { 290 @SuppressWarnings("unchecked") 291 E e = (E) element; 292 if (!range.contains(e) || root == null) { 293 return 0; 294 } 295 newRoot = root.remove(comparator(), e, occurrences, result); 296 } catch (ClassCastException | NullPointerException e) { 297 return 0; 298 } 299 rootReference.checkAndSet(root, newRoot); 300 return result[0]; 301 } 302 303 @CanIgnoreReturnValue 304 @Override 305 public int setCount(@ParametricNullness E element, int count) { 306 checkNonnegative(count, "count"); 307 if (!range.contains(element)) { 308 checkArgument(count == 0); 309 return 0; 310 } 311 312 AvlNode<E> root = rootReference.get(); 313 if (root == null) { 314 if (count > 0) { 315 add(element, count); 316 } 317 return 0; 318 } 319 int[] result = new int[1]; // used as a mutable int reference to hold result 320 AvlNode<E> newRoot = root.setCount(comparator(), element, count, result); 321 rootReference.checkAndSet(root, newRoot); 322 return result[0]; 323 } 324 325 @CanIgnoreReturnValue 326 @Override 327 public boolean setCount(@ParametricNullness E element, int oldCount, int newCount) { 328 checkNonnegative(newCount, "newCount"); 329 checkNonnegative(oldCount, "oldCount"); 330 checkArgument(range.contains(element)); 331 332 AvlNode<E> root = rootReference.get(); 333 if (root == null) { 334 if (oldCount == 0) { 335 if (newCount > 0) { 336 add(element, newCount); 337 } 338 return true; 339 } else { 340 return false; 341 } 342 } 343 int[] result = new int[1]; // used as a mutable int reference to hold result 344 AvlNode<E> newRoot = root.setCount(comparator(), element, oldCount, newCount, result); 345 rootReference.checkAndSet(root, newRoot); 346 return result[0] == oldCount; 347 } 348 349 @Override 350 public void clear() { 351 if (!range.hasLowerBound() && !range.hasUpperBound()) { 352 // We can do this in O(n) rather than removing one by one, which could force rebalancing. 353 for (AvlNode<E> current = header.succ(); current != header; ) { 354 AvlNode<E> next = current.succ(); 355 356 current.elemCount = 0; 357 // Also clear these fields so that one deleted Entry doesn't retain all elements. 358 current.left = null; 359 current.right = null; 360 current.pred = null; 361 current.succ = null; 362 363 current = next; 364 } 365 successor(header, header); 366 rootReference.clear(); 367 } else { 368 // TODO(cpovirk): Perhaps we can optimize in this case, too? 369 Iterators.clear(entryIterator()); 370 } 371 } 372 373 private Entry<E> wrapEntry(final AvlNode<E> baseEntry) { 374 return new Multisets.AbstractEntry<E>() { 375 @Override 376 @ParametricNullness 377 public E getElement() { 378 return baseEntry.getElement(); 379 } 380 381 @Override 382 public int getCount() { 383 int result = baseEntry.getCount(); 384 if (result == 0) { 385 return count(getElement()); 386 } else { 387 return result; 388 } 389 } 390 }; 391 } 392 393 /** Returns the first node in the tree that is in range. */ 394 @CheckForNull 395 private AvlNode<E> firstNode() { 396 AvlNode<E> root = rootReference.get(); 397 if (root == null) { 398 return null; 399 } 400 AvlNode<E> node; 401 if (range.hasLowerBound()) { 402 // The cast is safe because of the hasLowerBound check. 403 E endpoint = uncheckedCastNullableTToT(range.getLowerEndpoint()); 404 node = root.ceiling(comparator(), endpoint); 405 if (node == null) { 406 return null; 407 } 408 if (range.getLowerBoundType() == BoundType.OPEN 409 && comparator().compare(endpoint, node.getElement()) == 0) { 410 node = node.succ(); 411 } 412 } else { 413 node = header.succ(); 414 } 415 return (node == header || !range.contains(node.getElement())) ? null : node; 416 } 417 418 @CheckForNull 419 private AvlNode<E> lastNode() { 420 AvlNode<E> root = rootReference.get(); 421 if (root == null) { 422 return null; 423 } 424 AvlNode<E> node; 425 if (range.hasUpperBound()) { 426 // The cast is safe because of the hasUpperBound check. 427 E endpoint = uncheckedCastNullableTToT(range.getUpperEndpoint()); 428 node = root.floor(comparator(), endpoint); 429 if (node == null) { 430 return null; 431 } 432 if (range.getUpperBoundType() == BoundType.OPEN 433 && comparator().compare(endpoint, node.getElement()) == 0) { 434 node = node.pred(); 435 } 436 } else { 437 node = header.pred(); 438 } 439 return (node == header || !range.contains(node.getElement())) ? null : node; 440 } 441 442 @Override 443 Iterator<E> elementIterator() { 444 return Multisets.elementIterator(entryIterator()); 445 } 446 447 @Override 448 Iterator<Entry<E>> entryIterator() { 449 return new Iterator<Entry<E>>() { 450 @CheckForNull AvlNode<E> current = firstNode(); 451 @CheckForNull Entry<E> prevEntry; 452 453 @Override 454 public boolean hasNext() { 455 if (current == null) { 456 return false; 457 } else if (range.tooHigh(current.getElement())) { 458 current = null; 459 return false; 460 } else { 461 return true; 462 } 463 } 464 465 @Override 466 public Entry<E> next() { 467 if (!hasNext()) { 468 throw new NoSuchElementException(); 469 } 470 // requireNonNull is safe because current is only nulled out after iteration is complete. 471 Entry<E> result = wrapEntry(requireNonNull(current)); 472 prevEntry = result; 473 if (current.succ() == header) { 474 current = null; 475 } else { 476 current = current.succ(); 477 } 478 return result; 479 } 480 481 @Override 482 public void remove() { 483 checkState(prevEntry != null, "no calls to next() since the last call to remove()"); 484 setCount(prevEntry.getElement(), 0); 485 prevEntry = null; 486 } 487 }; 488 } 489 490 @Override 491 Iterator<Entry<E>> descendingEntryIterator() { 492 return new Iterator<Entry<E>>() { 493 @CheckForNull AvlNode<E> current = lastNode(); 494 @CheckForNull Entry<E> prevEntry = null; 495 496 @Override 497 public boolean hasNext() { 498 if (current == null) { 499 return false; 500 } else if (range.tooLow(current.getElement())) { 501 current = null; 502 return false; 503 } else { 504 return true; 505 } 506 } 507 508 @Override 509 public Entry<E> next() { 510 if (!hasNext()) { 511 throw new NoSuchElementException(); 512 } 513 // requireNonNull is safe because current is only nulled out after iteration is complete. 514 requireNonNull(current); 515 Entry<E> result = wrapEntry(current); 516 prevEntry = result; 517 if (current.pred() == header) { 518 current = null; 519 } else { 520 current = current.pred(); 521 } 522 return result; 523 } 524 525 @Override 526 public void remove() { 527 checkState(prevEntry != null, "no calls to next() since the last call to remove()"); 528 setCount(prevEntry.getElement(), 0); 529 prevEntry = null; 530 } 531 }; 532 } 533 534 @Override 535 public Iterator<E> iterator() { 536 return Multisets.iteratorImpl(this); 537 } 538 539 @Override 540 public SortedMultiset<E> headMultiset(@ParametricNullness E upperBound, BoundType boundType) { 541 return new TreeMultiset<>( 542 rootReference, 543 range.intersect(GeneralRange.upTo(comparator(), upperBound, boundType)), 544 header); 545 } 546 547 @Override 548 public SortedMultiset<E> tailMultiset(@ParametricNullness E lowerBound, BoundType boundType) { 549 return new TreeMultiset<>( 550 rootReference, 551 range.intersect(GeneralRange.downTo(comparator(), lowerBound, boundType)), 552 header); 553 } 554 555 private static final class Reference<T> { 556 @CheckForNull private T value; 557 558 @CheckForNull 559 public T get() { 560 return value; 561 } 562 563 public void checkAndSet(@CheckForNull T expected, @CheckForNull T newValue) { 564 if (value != expected) { 565 throw new ConcurrentModificationException(); 566 } 567 value = newValue; 568 } 569 570 void clear() { 571 value = null; 572 } 573 } 574 575 private static final class AvlNode<E extends @Nullable Object> { 576 /* 577 * For "normal" nodes, the type of this field is `E`, not `@Nullable E` (though note that E is a 578 * type that can include null, as in a TreeMultiset<@Nullable String>). 579 * 580 * For the header node, though, this field contains `null`, regardless of the type of the 581 * multiset. 582 * 583 * Most code that operates on an AvlNode never operates on the header node. Such code can access 584 * the elem field without a null check by calling getElement(). 585 */ 586 @CheckForNull private final E elem; 587 588 // elemCount is 0 iff this node has been deleted. 589 private int elemCount; 590 591 private int distinctElements; 592 private long totalCount; 593 private int height; 594 @CheckForNull private AvlNode<E> left; 595 @CheckForNull private AvlNode<E> right; 596 /* 597 * pred and succ are nullable after construction, but we always call successor() to initialize 598 * them immediately thereafter. 599 * 600 * They may be subsequently nulled out by TreeMultiset.clear(). I think that the only place that 601 * we can reference a node whose fields have been cleared is inside the iterator (and presumably 602 * only under concurrent modification). 603 * 604 * To access these fields when you know that they are not null, call the pred() and succ() 605 * methods, which perform null checks before returning the fields. 606 */ 607 @CheckForNull private AvlNode<E> pred; 608 @CheckForNull private AvlNode<E> succ; 609 610 AvlNode(@ParametricNullness E elem, int elemCount) { 611 checkArgument(elemCount > 0); 612 this.elem = elem; 613 this.elemCount = elemCount; 614 this.totalCount = elemCount; 615 this.distinctElements = 1; 616 this.height = 1; 617 this.left = null; 618 this.right = null; 619 } 620 621 /** Constructor for the header node. */ 622 AvlNode() { 623 this.elem = null; 624 this.elemCount = 1; 625 } 626 627 // For discussion of pred() and succ(), see the comment on the pred and succ fields. 628 629 private AvlNode<E> pred() { 630 return requireNonNull(pred); 631 } 632 633 private AvlNode<E> succ() { 634 return requireNonNull(succ); 635 } 636 637 int count(Comparator<? super E> comparator, @ParametricNullness E e) { 638 int cmp = comparator.compare(e, getElement()); 639 if (cmp < 0) { 640 return (left == null) ? 0 : left.count(comparator, e); 641 } else if (cmp > 0) { 642 return (right == null) ? 0 : right.count(comparator, e); 643 } else { 644 return elemCount; 645 } 646 } 647 648 private AvlNode<E> addRightChild(@ParametricNullness E e, int count) { 649 right = new AvlNode<>(e, count); 650 successor(this, right, succ()); 651 height = max(2, height); 652 distinctElements++; 653 totalCount += count; 654 return this; 655 } 656 657 private AvlNode<E> addLeftChild(@ParametricNullness E e, int count) { 658 left = new AvlNode<>(e, count); 659 successor(pred(), left, this); 660 height = max(2, height); 661 distinctElements++; 662 totalCount += count; 663 return this; 664 } 665 666 AvlNode<E> add( 667 Comparator<? super E> comparator, @ParametricNullness E e, int count, int[] result) { 668 /* 669 * It speeds things up considerably to unconditionally add count to totalCount here, 670 * but that destroys failure atomicity in the case of count overflow. =( 671 */ 672 int cmp = comparator.compare(e, getElement()); 673 if (cmp < 0) { 674 AvlNode<E> initLeft = left; 675 if (initLeft == null) { 676 result[0] = 0; 677 return addLeftChild(e, count); 678 } 679 int initHeight = initLeft.height; 680 681 left = initLeft.add(comparator, e, count, result); 682 if (result[0] == 0) { 683 distinctElements++; 684 } 685 this.totalCount += count; 686 return (left.height == initHeight) ? this : rebalance(); 687 } else if (cmp > 0) { 688 AvlNode<E> initRight = right; 689 if (initRight == null) { 690 result[0] = 0; 691 return addRightChild(e, count); 692 } 693 int initHeight = initRight.height; 694 695 right = initRight.add(comparator, e, count, result); 696 if (result[0] == 0) { 697 distinctElements++; 698 } 699 this.totalCount += count; 700 return (right.height == initHeight) ? this : rebalance(); 701 } 702 703 // adding count to me! No rebalance possible. 704 result[0] = elemCount; 705 long resultCount = (long) elemCount + count; 706 checkArgument(resultCount <= Integer.MAX_VALUE); 707 this.elemCount += count; 708 this.totalCount += count; 709 return this; 710 } 711 712 @CheckForNull 713 AvlNode<E> remove( 714 Comparator<? super E> comparator, @ParametricNullness E e, int count, int[] result) { 715 int cmp = comparator.compare(e, getElement()); 716 if (cmp < 0) { 717 AvlNode<E> initLeft = left; 718 if (initLeft == null) { 719 result[0] = 0; 720 return this; 721 } 722 723 left = initLeft.remove(comparator, e, count, result); 724 725 if (result[0] > 0) { 726 if (count >= result[0]) { 727 this.distinctElements--; 728 this.totalCount -= result[0]; 729 } else { 730 this.totalCount -= count; 731 } 732 } 733 return (result[0] == 0) ? this : rebalance(); 734 } else if (cmp > 0) { 735 AvlNode<E> initRight = right; 736 if (initRight == null) { 737 result[0] = 0; 738 return this; 739 } 740 741 right = initRight.remove(comparator, e, count, result); 742 743 if (result[0] > 0) { 744 if (count >= result[0]) { 745 this.distinctElements--; 746 this.totalCount -= result[0]; 747 } else { 748 this.totalCount -= count; 749 } 750 } 751 return rebalance(); 752 } 753 754 // removing count from me! 755 result[0] = elemCount; 756 if (count >= elemCount) { 757 return deleteMe(); 758 } else { 759 this.elemCount -= count; 760 this.totalCount -= count; 761 return this; 762 } 763 } 764 765 @CheckForNull 766 AvlNode<E> setCount( 767 Comparator<? super E> comparator, @ParametricNullness E e, int count, int[] result) { 768 int cmp = comparator.compare(e, getElement()); 769 if (cmp < 0) { 770 AvlNode<E> initLeft = left; 771 if (initLeft == null) { 772 result[0] = 0; 773 return (count > 0) ? addLeftChild(e, count) : this; 774 } 775 776 left = initLeft.setCount(comparator, e, count, result); 777 778 if (count == 0 && result[0] != 0) { 779 this.distinctElements--; 780 } else if (count > 0 && result[0] == 0) { 781 this.distinctElements++; 782 } 783 784 this.totalCount += count - result[0]; 785 return rebalance(); 786 } else if (cmp > 0) { 787 AvlNode<E> initRight = right; 788 if (initRight == null) { 789 result[0] = 0; 790 return (count > 0) ? addRightChild(e, count) : this; 791 } 792 793 right = initRight.setCount(comparator, e, count, result); 794 795 if (count == 0 && result[0] != 0) { 796 this.distinctElements--; 797 } else if (count > 0 && result[0] == 0) { 798 this.distinctElements++; 799 } 800 801 this.totalCount += count - result[0]; 802 return rebalance(); 803 } 804 805 // setting my count 806 result[0] = elemCount; 807 if (count == 0) { 808 return deleteMe(); 809 } 810 this.totalCount += count - elemCount; 811 this.elemCount = count; 812 return this; 813 } 814 815 @CheckForNull 816 AvlNode<E> setCount( 817 Comparator<? super E> comparator, 818 @ParametricNullness E e, 819 int expectedCount, 820 int newCount, 821 int[] result) { 822 int cmp = comparator.compare(e, getElement()); 823 if (cmp < 0) { 824 AvlNode<E> initLeft = left; 825 if (initLeft == null) { 826 result[0] = 0; 827 if (expectedCount == 0 && newCount > 0) { 828 return addLeftChild(e, newCount); 829 } 830 return this; 831 } 832 833 left = initLeft.setCount(comparator, e, expectedCount, newCount, result); 834 835 if (result[0] == expectedCount) { 836 if (newCount == 0 && result[0] != 0) { 837 this.distinctElements--; 838 } else if (newCount > 0 && result[0] == 0) { 839 this.distinctElements++; 840 } 841 this.totalCount += newCount - result[0]; 842 } 843 return rebalance(); 844 } else if (cmp > 0) { 845 AvlNode<E> initRight = right; 846 if (initRight == null) { 847 result[0] = 0; 848 if (expectedCount == 0 && newCount > 0) { 849 return addRightChild(e, newCount); 850 } 851 return this; 852 } 853 854 right = initRight.setCount(comparator, e, expectedCount, newCount, result); 855 856 if (result[0] == expectedCount) { 857 if (newCount == 0 && result[0] != 0) { 858 this.distinctElements--; 859 } else if (newCount > 0 && result[0] == 0) { 860 this.distinctElements++; 861 } 862 this.totalCount += newCount - result[0]; 863 } 864 return rebalance(); 865 } 866 867 // setting my count 868 result[0] = elemCount; 869 if (expectedCount == elemCount) { 870 if (newCount == 0) { 871 return deleteMe(); 872 } 873 this.totalCount += newCount - elemCount; 874 this.elemCount = newCount; 875 } 876 return this; 877 } 878 879 @CheckForNull 880 private AvlNode<E> deleteMe() { 881 int oldElemCount = this.elemCount; 882 this.elemCount = 0; 883 successor(pred(), succ()); 884 if (left == null) { 885 return right; 886 } else if (right == null) { 887 return left; 888 } else if (left.height >= right.height) { 889 AvlNode<E> newTop = pred(); 890 // newTop is the maximum node in my left subtree 891 newTop.left = left.removeMax(newTop); 892 newTop.right = right; 893 newTop.distinctElements = distinctElements - 1; 894 newTop.totalCount = totalCount - oldElemCount; 895 return newTop.rebalance(); 896 } else { 897 AvlNode<E> newTop = succ(); 898 newTop.right = right.removeMin(newTop); 899 newTop.left = left; 900 newTop.distinctElements = distinctElements - 1; 901 newTop.totalCount = totalCount - oldElemCount; 902 return newTop.rebalance(); 903 } 904 } 905 906 // Removes the minimum node from this subtree to be reused elsewhere 907 @CheckForNull 908 private AvlNode<E> removeMin(AvlNode<E> node) { 909 if (left == null) { 910 return right; 911 } else { 912 left = left.removeMin(node); 913 distinctElements--; 914 totalCount -= node.elemCount; 915 return rebalance(); 916 } 917 } 918 919 // Removes the maximum node from this subtree to be reused elsewhere 920 @CheckForNull 921 private AvlNode<E> removeMax(AvlNode<E> node) { 922 if (right == null) { 923 return left; 924 } else { 925 right = right.removeMax(node); 926 distinctElements--; 927 totalCount -= node.elemCount; 928 return rebalance(); 929 } 930 } 931 932 private void recomputeMultiset() { 933 this.distinctElements = 934 1 + TreeMultiset.distinctElements(left) + TreeMultiset.distinctElements(right); 935 this.totalCount = elemCount + totalCount(left) + totalCount(right); 936 } 937 938 private void recomputeHeight() { 939 this.height = 1 + max(height(left), height(right)); 940 } 941 942 private void recompute() { 943 recomputeMultiset(); 944 recomputeHeight(); 945 } 946 947 private AvlNode<E> rebalance() { 948 switch (balanceFactor()) { 949 case -2: 950 // requireNonNull is safe because right must exist in order to get a negative factor. 951 requireNonNull(right); 952 if (right.balanceFactor() > 0) { 953 right = right.rotateRight(); 954 } 955 return rotateLeft(); 956 case 2: 957 // requireNonNull is safe because left must exist in order to get a positive factor. 958 requireNonNull(left); 959 if (left.balanceFactor() < 0) { 960 left = left.rotateLeft(); 961 } 962 return rotateRight(); 963 default: 964 recomputeHeight(); 965 return this; 966 } 967 } 968 969 private int balanceFactor() { 970 return height(left) - height(right); 971 } 972 973 private AvlNode<E> rotateLeft() { 974 checkState(right != null); 975 AvlNode<E> newTop = right; 976 this.right = newTop.left; 977 newTop.left = this; 978 newTop.totalCount = this.totalCount; 979 newTop.distinctElements = this.distinctElements; 980 this.recompute(); 981 newTop.recomputeHeight(); 982 return newTop; 983 } 984 985 private AvlNode<E> rotateRight() { 986 checkState(left != null); 987 AvlNode<E> newTop = left; 988 this.left = newTop.right; 989 newTop.right = this; 990 newTop.totalCount = this.totalCount; 991 newTop.distinctElements = this.distinctElements; 992 this.recompute(); 993 newTop.recomputeHeight(); 994 return newTop; 995 } 996 997 private static long totalCount(@CheckForNull AvlNode<?> node) { 998 return (node == null) ? 0 : node.totalCount; 999 } 1000 1001 private static int height(@CheckForNull AvlNode<?> node) { 1002 return (node == null) ? 0 : node.height; 1003 } 1004 1005 @CheckForNull 1006 private AvlNode<E> ceiling(Comparator<? super E> comparator, @ParametricNullness E e) { 1007 int cmp = comparator.compare(e, getElement()); 1008 if (cmp < 0) { 1009 return (left == null) ? this : MoreObjects.firstNonNull(left.ceiling(comparator, e), this); 1010 } else if (cmp == 0) { 1011 return this; 1012 } else { 1013 return (right == null) ? null : right.ceiling(comparator, e); 1014 } 1015 } 1016 1017 @CheckForNull 1018 private AvlNode<E> floor(Comparator<? super E> comparator, @ParametricNullness E e) { 1019 int cmp = comparator.compare(e, getElement()); 1020 if (cmp > 0) { 1021 return (right == null) ? this : MoreObjects.firstNonNull(right.floor(comparator, e), this); 1022 } else if (cmp == 0) { 1023 return this; 1024 } else { 1025 return (left == null) ? null : left.floor(comparator, e); 1026 } 1027 } 1028 1029 @ParametricNullness 1030 E getElement() { 1031 // For discussion of this cast, see the comment on the elem field. 1032 return uncheckedCastNullableTToT(elem); 1033 } 1034 1035 int getCount() { 1036 return elemCount; 1037 } 1038 1039 @Override 1040 public String toString() { 1041 return Multisets.immutableEntry(getElement(), getCount()).toString(); 1042 } 1043 } 1044 1045 private static <T extends @Nullable Object> void successor(AvlNode<T> a, AvlNode<T> b) { 1046 a.succ = b; 1047 b.pred = a; 1048 } 1049 1050 private static <T extends @Nullable Object> void successor( 1051 AvlNode<T> a, AvlNode<T> b, AvlNode<T> c) { 1052 successor(a, b); 1053 successor(b, c); 1054 } 1055 1056 /* 1057 * TODO(jlevy): Decide whether entrySet() should return entries with an equals() method that 1058 * calls the comparator to compare the two keys. If that change is made, 1059 * AbstractMultiset.equals() can simply check whether two multisets have equal entry sets. 1060 */ 1061 1062 /** 1063 * @serialData the comparator, the number of distinct elements, the first element, its count, the 1064 * second element, its count, and so on 1065 */ 1066 @J2ktIncompatible 1067 @GwtIncompatible // java.io.ObjectOutputStream 1068 private void writeObject(ObjectOutputStream stream) throws IOException { 1069 stream.defaultWriteObject(); 1070 stream.writeObject(elementSet().comparator()); 1071 Serialization.writeMultiset(this, stream); 1072 } 1073 1074 @J2ktIncompatible 1075 @GwtIncompatible // java.io.ObjectInputStream 1076 private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException { 1077 stream.defaultReadObject(); 1078 @SuppressWarnings("unchecked") 1079 // reading data stored by writeObject 1080 Comparator<? super E> comparator = (Comparator<? super E>) requireNonNull(stream.readObject()); 1081 Serialization.getFieldSetter(AbstractSortedMultiset.class, "comparator").set(this, comparator); 1082 Serialization.getFieldSetter(TreeMultiset.class, "range") 1083 .set(this, GeneralRange.all(comparator)); 1084 Serialization.getFieldSetter(TreeMultiset.class, "rootReference") 1085 .set(this, new Reference<AvlNode<E>>()); 1086 AvlNode<E> header = new AvlNode<>(); 1087 Serialization.getFieldSetter(TreeMultiset.class, "header").set(this, header); 1088 successor(header, header); 1089 Serialization.populateMultiset(this, stream); 1090 } 1091 1092 @GwtIncompatible // not needed in emulated source 1093 @J2ktIncompatible 1094 private static final long serialVersionUID = 1; 1095}