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