1/*
2 * Copyright (C) 2011 The Guava Authors
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
5 * in compliance with the License. You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software distributed under the
10 * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
11 * express or implied. See the License for the specific language governing permissions and
12 * limitations under the License.
13 */
14
15package com.google.common.collect;
16
17import static com.google.common.base.Preconditions.checkNotNull;
18import static com.google.common.collect.BstOperations.extractMax;
19import static com.google.common.collect.BstOperations.extractMin;
20import static com.google.common.collect.BstOperations.insertMax;
21import static com.google.common.collect.BstOperations.insertMin;
22import static com.google.common.collect.BstSide.LEFT;
23import static com.google.common.collect.BstSide.RIGHT;
24
25import com.google.common.annotations.GwtCompatible;
26
27import javax.annotation.Nullable;
28
29/**
30 * A tree-size-based set of balancing policies, based on <a
31 * href="http://www.swiss.ai.mit.edu/~adams/BB/"> Stephen Adams, "Efficient sets: a balancing
32 * act."</a>.
33 *
34 * @author Louis Wasserman
35 */
36@GwtCompatible
37final class BstCountBasedBalancePolicies {
38  private BstCountBasedBalancePolicies() {}
39
40  private static final int SINGLE_ROTATE_RATIO = 4;
41  private static final int SECOND_ROTATE_RATIO = 2;
42
43  /**
44   * Returns a balance policy that does no balancing or the bare minimum (for {@code combine}).
45   */
46  public static <N extends BstNode<?, N>> BstBalancePolicy<N> noRebalancePolicy(
47      final BstAggregate<N> countAggregate) {
48    checkNotNull(countAggregate);
49    return new BstBalancePolicy<N>() {
50      @Override
51      public N balance(
52          BstNodeFactory<N> nodeFactory, N source, @Nullable N left, @Nullable N right) {
53        return checkNotNull(nodeFactory).createNode(source, left, right);
54      }
55
56      @Nullable
57      @Override
58      public N combine(BstNodeFactory<N> nodeFactory, @Nullable N left, @Nullable N right) {
59        if (left == null) {
60          return right;
61        } else if (right == null) {
62          return left;
63        } else if (countAggregate.treeValue(left) > countAggregate.treeValue(right)) {
64          return nodeFactory.createNode(
65              left, left.childOrNull(LEFT), combine(nodeFactory, left.childOrNull(RIGHT), right));
66        } else {
67          return nodeFactory.createNode(right, combine(nodeFactory, left, right.childOrNull(LEFT)),
68              right.childOrNull(RIGHT));
69        }
70      }
71    };
72  }
73
74  /**
75   * Returns a balance policy that expects the sizes of each side to be at most one node (added or
76   * removed) away from being balanced. {@code balance} takes {@code O(1)} time, and {@code
77   * combine} takes {@code O(log n)} time.
78   */
79  public static <K, N extends BstNode<K, N>> BstBalancePolicy<N> singleRebalancePolicy(
80      final BstAggregate<N> countAggregate) {
81    checkNotNull(countAggregate);
82    return new BstBalancePolicy<N>() {
83      @Override
84      public N balance(
85          BstNodeFactory<N> nodeFactory, N source, @Nullable N left, @Nullable N right) {
86        long countL = countAggregate.treeValue(left);
87        long countR = countAggregate.treeValue(right);
88        if (countL + countR > 1) {
89          if (countR >= SINGLE_ROTATE_RATIO * countL) {
90            return rotateL(nodeFactory, source, left, right);
91          } else if (countL >= SINGLE_ROTATE_RATIO * countR) {
92            return rotateR(nodeFactory, source, left, right);
93          }
94        }
95        return nodeFactory.createNode(source, left, right);
96      }
97
98      private N rotateL(BstNodeFactory<N> nodeFactory, N source, @Nullable N left, N right) {
99        checkNotNull(right);
100        N rl = right.childOrNull(LEFT);
101        N rr = right.childOrNull(RIGHT);
102        if (countAggregate.treeValue(rl) >= SECOND_ROTATE_RATIO * countAggregate.treeValue(rr)) {
103          right = singleR(nodeFactory, right, rl, rr);
104        }
105        return singleL(nodeFactory, source, left, right);
106      }
107
108      private N rotateR(BstNodeFactory<N> nodeFactory, N source, N left, @Nullable N right) {
109        checkNotNull(left);
110        N lr = left.childOrNull(RIGHT);
111        N ll = left.childOrNull(LEFT);
112        if (countAggregate.treeValue(lr) >= SECOND_ROTATE_RATIO * countAggregate.treeValue(ll)) {
113          left = singleL(nodeFactory, left, ll, lr);
114        }
115        return singleR(nodeFactory, source, left, right);
116      }
117
118      private N singleL(BstNodeFactory<N> nodeFactory, N source, @Nullable N left, N right) {
119        checkNotNull(right);
120        return nodeFactory.createNode(right,
121            nodeFactory.createNode(source, left, right.childOrNull(LEFT)),
122            right.childOrNull(RIGHT));
123      }
124
125      private N singleR(BstNodeFactory<N> nodeFactory, N source, N left, @Nullable N right) {
126        checkNotNull(left);
127        return nodeFactory.createNode(left, left.childOrNull(LEFT),
128            nodeFactory.createNode(source, left.childOrNull(RIGHT), right));
129      }
130
131      @Nullable
132      @Override
133      public N combine(BstNodeFactory<N> nodeFactory, @Nullable N left, @Nullable N right) {
134        if (left == null) {
135          return right;
136        } else if (right == null) {
137          return left;
138        }
139        N newRootSource;
140        if (countAggregate.treeValue(left) > countAggregate.treeValue(right)) {
141          BstMutationResult<K, N> extractLeftMax = extractMax(left, nodeFactory, this);
142          newRootSource = extractLeftMax.getOriginalTarget();
143          left = extractLeftMax.getChangedRoot();
144        } else {
145          BstMutationResult<K, N> extractRightMin = extractMin(right, nodeFactory, this);
146          newRootSource = extractRightMin.getOriginalTarget();
147          right = extractRightMin.getChangedRoot();
148        }
149        return nodeFactory.createNode(newRootSource, left, right);
150      }
151    };
152  }
153
154  /**
155   * Returns a balance policy that makes no assumptions on the relative balance of the two sides
156   * and performs a full rebalancing as necessary. Both {@code balance} and {@code combine} take
157   * {@code O(log n)} time.
158   */
159  public static <K, N extends BstNode<K, N>> BstBalancePolicy<N> fullRebalancePolicy(
160      final BstAggregate<N> countAggregate) {
161    checkNotNull(countAggregate);
162    final BstBalancePolicy<N> singleBalancePolicy =
163        BstCountBasedBalancePolicies.<K, N>singleRebalancePolicy(countAggregate);
164    return new BstBalancePolicy<N>() {
165      @Override
166      public N balance(
167          BstNodeFactory<N> nodeFactory, N source, @Nullable N left, @Nullable N right) {
168        if (left == null) {
169          return insertMin(right, source, nodeFactory, singleBalancePolicy);
170        } else if (right == null) {
171          return insertMax(left, source, nodeFactory, singleBalancePolicy);
172        }
173        long countL = countAggregate.treeValue(left);
174        long countR = countAggregate.treeValue(right);
175        if (SINGLE_ROTATE_RATIO * countL <= countR) {
176          N resultLeft = balance(nodeFactory, source, left, right.childOrNull(LEFT));
177          return singleBalancePolicy.balance(
178              nodeFactory, right, resultLeft, right.childOrNull(RIGHT));
179        } else if (SINGLE_ROTATE_RATIO * countR <= countL) {
180          N resultRight = balance(nodeFactory, source, left.childOrNull(RIGHT), right);
181          return singleBalancePolicy.balance(
182              nodeFactory, left, left.childOrNull(LEFT), resultRight);
183        } else {
184          return nodeFactory.createNode(source, left, right);
185        }
186      }
187
188      @Nullable
189      @Override
190      public N combine(BstNodeFactory<N> nodeFactory, @Nullable N left, @Nullable N right) {
191        if (left == null) {
192          return right;
193        } else if (right == null) {
194          return left;
195        }
196        long countL = countAggregate.treeValue(left);
197        long countR = countAggregate.treeValue(right);
198        if (SINGLE_ROTATE_RATIO * countL <= countR) {
199          N resultLeft = combine(nodeFactory, left, right.childOrNull(LEFT));
200          return singleBalancePolicy.balance(
201              nodeFactory, right, resultLeft, right.childOrNull(RIGHT));
202        } else if (SINGLE_ROTATE_RATIO * countR <= countL) {
203          N resultRight = combine(nodeFactory, left.childOrNull(RIGHT), right);
204          return singleBalancePolicy.balance(
205              nodeFactory, left, left.childOrNull(LEFT), resultRight);
206        } else {
207          return singleBalancePolicy.combine(nodeFactory, left, right);
208        }
209      }
210    };
211  }
212}
213