View Javadoc
1   /*
2    * Copyright (C) 2012 The Guava Authors
3    *
4    * Licensed under the Apache License, Version 2.0 (the "License");
5    * you may not use this file except in compliance with the License.
6    * You may obtain a copy of the License at
7    *
8    * http://www.apache.org/licenses/LICENSE-2.0
9    *
10   * Unless required by applicable law or agreed to in writing, software
11   * distributed under the License is distributed on an "AS IS" BASIS,
12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13   * See the License for the specific language governing permissions and
14   * limitations under the License.
15   */
16  
17  package com.google.common.collect;
18  
19  import static com.google.common.base.Preconditions.checkNotNull;
20  
21  import com.google.common.annotations.Beta;
22  import com.google.common.annotations.GwtCompatible;
23  import com.google.common.base.Optional;
24  import java.util.ArrayDeque;
25  import java.util.BitSet;
26  import java.util.Deque;
27  import java.util.Iterator;
28  import java.util.function.Consumer;
29  
30  /**
31   * A variant of {@link TreeTraverser} for binary trees, providing additional traversals specific to
32   * binary trees.
33   *
34   * @author Louis Wasserman
35   * @since 15.0
36   */
37  @Beta
38  @GwtCompatible
39  public abstract class BinaryTreeTraverser<T> extends TreeTraverser<T> {
40  
41    /**
42     * Returns the left child of the specified node, or {@link Optional#absent()} if the specified
43     * node has no left child.
44     */
45    public abstract Optional<T> leftChild(T root);
46  
47    /**
48     * Returns the right child of the specified node, or {@link Optional#absent()} if the specified
49     * node has no right child.
50     */
51    public abstract Optional<T> rightChild(T root);
52  
53    /**
54     * Returns the children of this node, in left-to-right order.
55     */
56    @Override
57    public final Iterable<T> children(final T root) {
58      checkNotNull(root);
59      return new FluentIterable<T>() {
60        @Override
61        public Iterator<T> iterator() {
62          return new AbstractIterator<T>() {
63            boolean doneLeft;
64            boolean doneRight;
65  
66            @Override
67            protected T computeNext() {
68              if (!doneLeft) {
69                doneLeft = true;
70                Optional<T> left = leftChild(root);
71                if (left.isPresent()) {
72                  return left.get();
73                }
74              }
75              if (!doneRight) {
76                doneRight = true;
77                Optional<T> right = rightChild(root);
78                if (right.isPresent()) {
79                  return right.get();
80                }
81              }
82              return endOfData();
83            }
84          };
85        }
86  
87        @Override
88        public void forEach(Consumer<? super T> action) {
89          acceptIfPresent(action, leftChild(root));
90          acceptIfPresent(action, rightChild(root));
91        }
92      };
93    }
94  
95    @Override
96    UnmodifiableIterator<T> preOrderIterator(T root) {
97      return new PreOrderIterator(root);
98    }
99  
100   /*
101    * Optimized implementation of preOrderIterator for binary trees.
102    */
103   private final class PreOrderIterator extends UnmodifiableIterator<T>
104       implements PeekingIterator<T> {
105     private final Deque<T> stack;
106 
107     PreOrderIterator(T root) {
108       this.stack = new ArrayDeque<T>(8);
109       stack.addLast(root);
110     }
111 
112     @Override
113     public boolean hasNext() {
114       return !stack.isEmpty();
115     }
116 
117     @Override
118     public T next() {
119       T result = stack.removeLast();
120       pushIfPresent(stack, rightChild(result));
121       pushIfPresent(stack, leftChild(result));
122       return result;
123     }
124 
125     @Override
126     public T peek() {
127       return stack.getLast();
128     }
129   }
130 
131   @Override
132   UnmodifiableIterator<T> postOrderIterator(T root) {
133     return new PostOrderIterator(root);
134   }
135 
136   /*
137    * Optimized implementation of postOrderIterator for binary trees.
138    */
139   private final class PostOrderIterator extends UnmodifiableIterator<T> {
140     private final Deque<T> stack;
141     private final BitSet hasExpanded;
142 
143     PostOrderIterator(T root) {
144       this.stack = new ArrayDeque<T>(8);
145       stack.addLast(root);
146       this.hasExpanded = new BitSet();
147     }
148 
149     @Override
150     public boolean hasNext() {
151       return !stack.isEmpty();
152     }
153 
154     @Override
155     public T next() {
156       while (true) {
157         T node = stack.getLast();
158         boolean expandedNode = hasExpanded.get(stack.size() - 1);
159         if (expandedNode) {
160           stack.removeLast();
161           hasExpanded.clear(stack.size());
162           return node;
163         } else {
164           hasExpanded.set(stack.size() - 1);
165           pushIfPresent(stack, rightChild(node));
166           pushIfPresent(stack, leftChild(node));
167         }
168       }
169     }
170   }
171 
172   // TODO(lowasser): see if any significant optimizations are possible for breadthFirstIterator
173 
174   public final FluentIterable<T> inOrderTraversal(final T root) {
175     checkNotNull(root);
176     return new FluentIterable<T>() {
177       @Override
178       public UnmodifiableIterator<T> iterator() {
179         return new InOrderIterator(root);
180       }
181 
182       @Override
183       public void forEach(Consumer<? super T> action) {
184         checkNotNull(action);
185         new Consumer<T>() {
186           @Override
187           public void accept(T t) {
188             acceptIfPresent(this, leftChild(t));
189             action.accept(t);
190             acceptIfPresent(this, rightChild(t));
191           }
192         }.accept(root);
193       }
194     };
195   }
196 
197   private final class InOrderIterator extends AbstractIterator<T> {
198     private final Deque<T> stack;
199     private final BitSet hasExpandedLeft;
200 
201     InOrderIterator(T root) {
202       this.stack = new ArrayDeque<T>(8);
203       this.hasExpandedLeft = new BitSet();
204       stack.addLast(root);
205     }
206 
207     @Override
208     protected T computeNext() {
209       while (!stack.isEmpty()) {
210         T node = stack.getLast();
211         if (hasExpandedLeft.get(stack.size() - 1)) {
212           stack.removeLast();
213           hasExpandedLeft.clear(stack.size());
214           pushIfPresent(stack, rightChild(node));
215           return node;
216         } else {
217           hasExpandedLeft.set(stack.size() - 1);
218           pushIfPresent(stack, leftChild(node));
219         }
220       }
221       return endOfData();
222     }
223   }
224 
225   private static <T> void pushIfPresent(Deque<T> stack, Optional<T> node) {
226     if (node.isPresent()) {
227       stack.addLast(node.get());
228     }
229   }
230 
231   private static <T> void acceptIfPresent(Consumer<? super T> action, Optional<T> node) {
232     if (node.isPresent()) {
233       action.accept(node.get());
234     }
235   }
236 }