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.BstSide.LEFT;
19import static com.google.common.collect.BstSide.RIGHT;
20
21import com.google.common.annotations.GwtCompatible;
22
23import javax.annotation.Nullable;
24
25/**
26 * A utility class with operations on binary search trees that operate on some interval.
27 *
28 * @author Louis Wasserman
29 */
30@GwtCompatible
31final class BstRangeOps {
32  /**
33   * Returns the total value of the specified aggregation function on the specified tree restricted
34   * to the specified range. Assumes that the tree satisfies the binary search ordering property
35   * relative to {@code range.comparator()}.
36   */
37  public static <K, N extends BstNode<K, N>> long totalInRange(
38      BstAggregate<? super N> aggregate, GeneralRange<K> range, @Nullable N root) {
39    checkNotNull(aggregate);
40    checkNotNull(range);
41    if (root == null || range.isEmpty()) {
42      return 0;
43    }
44    long total = aggregate.treeValue(root);
45    if (range.hasLowerBound()) {
46      total -= totalBeyondRangeToSide(aggregate, range, LEFT, root);
47    }
48    if (range.hasUpperBound()) {
49      total -= totalBeyondRangeToSide(aggregate, range, RIGHT, root);
50    }
51    return total;
52  }
53
54  // Returns total value strictly to the specified side of the specified range.
55  private static <K, N extends BstNode<K, N>> long totalBeyondRangeToSide(
56      BstAggregate<? super N> aggregate, GeneralRange<K> range, BstSide side, @Nullable N root) {
57    long accum = 0;
58    while (root != null) {
59      if (beyond(range, root.getKey(), side)) {
60        accum += aggregate.entryValue(root);
61        accum += aggregate.treeValue(root.childOrNull(side));
62        root = root.childOrNull(side.other());
63      } else {
64        root = root.childOrNull(side);
65      }
66    }
67    return accum;
68  }
69
70  /**
71   * Returns a balanced tree containing all nodes from the specified tree that were <i>not</i> in
72   * the specified range, using the specified balance policy. Assumes that the tree satisfies the
73   * binary search ordering property relative to {@code range.comparator()}.
74   */
75  @Nullable
76  public static <K, N extends BstNode<K, N>> N minusRange(GeneralRange<K> range,
77      BstBalancePolicy<N> balancePolicy, BstNodeFactory<N> nodeFactory, @Nullable N root) {
78    checkNotNull(range);
79    checkNotNull(balancePolicy);
80    checkNotNull(nodeFactory);
81    N higher = range.hasUpperBound()
82        ? subTreeBeyondRangeToSide(range, balancePolicy, nodeFactory, RIGHT, root)
83        : null;
84    N lower = range.hasLowerBound()
85        ? subTreeBeyondRangeToSide(range, balancePolicy, nodeFactory, LEFT, root)
86        : null;
87    return balancePolicy.combine(nodeFactory, lower, higher);
88  }
89
90  /*
91   * Returns a balanced tree containing all nodes in the specified tree that are strictly to the
92   * specified side of the specified range.
93   */
94  @Nullable
95  private static <K, N extends BstNode<K, N>> N subTreeBeyondRangeToSide(GeneralRange<K> range,
96      BstBalancePolicy<N> balancePolicy, BstNodeFactory<N> nodeFactory, BstSide side,
97      @Nullable N root) {
98    if (root == null) {
99      return null;
100    }
101    if (beyond(range, root.getKey(), side)) {
102      N left = root.childOrNull(LEFT);
103      N right = root.childOrNull(RIGHT);
104      switch (side) {
105        case LEFT:
106          right = subTreeBeyondRangeToSide(range, balancePolicy, nodeFactory, LEFT, right);
107          break;
108        case RIGHT:
109          left = subTreeBeyondRangeToSide(range, balancePolicy, nodeFactory, RIGHT, left);
110          break;
111        default:
112          throw new AssertionError();
113      }
114      return balancePolicy.balance(nodeFactory, root, left, right);
115    } else {
116      return subTreeBeyondRangeToSide(
117          range, balancePolicy, nodeFactory, side, root.childOrNull(side));
118    }
119  }
120
121  /**
122   * Returns the furthest path to the specified side in the specified tree that falls into the
123   * specified range.
124   */
125  @Nullable
126  public static <K, N extends BstNode<K, N>, P extends BstPath<N, P>> P furthestPath(
127      GeneralRange<K> range, BstSide side, BstPathFactory<N, P> pathFactory, @Nullable N root) {
128    checkNotNull(range);
129    checkNotNull(pathFactory);
130    checkNotNull(side);
131    if (root == null) {
132      return null;
133    }
134    P path = pathFactory.initialPath(root);
135    return furthestPath(range, side, pathFactory, path);
136  }
137
138  private static <K, N extends BstNode<K, N>, P extends BstPath<N, P>> P furthestPath(
139      GeneralRange<K> range, BstSide side, BstPathFactory<N, P> pathFactory, P currentPath) {
140    N tip = currentPath.getTip();
141    K tipKey = tip.getKey();
142    if (beyond(range, tipKey, side)) {
143      if (tip.hasChild(side.other())) {
144        currentPath = pathFactory.extension(currentPath, side.other());
145        return furthestPath(range, side, pathFactory, currentPath);
146      } else {
147        return null;
148      }
149    } else if (tip.hasChild(side)) {
150      P alphaPath = pathFactory.extension(currentPath, side);
151      alphaPath = furthestPath(range, side, pathFactory, alphaPath);
152      if (alphaPath != null) {
153        return alphaPath;
154      }
155    }
156    return beyond(range, tipKey, side.other()) ? null : currentPath;
157  }
158
159  /**
160   * Returns {@code true} if {@code key} is beyond the specified side of the specified range.
161   */
162  public static <K> boolean beyond(GeneralRange<K> range, @Nullable K key, BstSide side) {
163    checkNotNull(range);
164    switch (side) {
165      case LEFT:
166        return range.tooLow(key);
167      case RIGHT:
168        return range.tooHigh(key);
169      default:
170        throw new AssertionError();
171    }
172  }
173
174  private BstRangeOps() {}
175}
176