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
17package com.google.common.collect;
18
19import static com.google.common.base.Preconditions.checkNotNull;
20
21import com.google.common.annotations.Beta;
22import com.google.common.annotations.GwtCompatible;
23import com.google.common.base.Optional;
24
25import java.util.ArrayDeque;
26import java.util.BitSet;
27import java.util.Deque;
28import java.util.Iterator;
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(emulated = true)
39public abstract class BinaryTreeTraverser<T> extends TreeTraverser<T> {
40  // TODO(user): make this GWT-compatible when we've checked in ArrayDeque and BitSet emulation
41
42  /**
43   * Returns the left child of the specified node, or {@link Optional#absent()} if the specified
44   * node has no left child.
45   */
46  public abstract Optional<T> leftChild(T root);
47
48  /**
49   * Returns the right child of the specified node, or {@link Optional#absent()} if the specified
50   * node has no right child.
51   */
52  public abstract Optional<T> rightChild(T root);
53
54  /**
55   * Returns the children of this node, in left-to-right order.
56   */
57  @Override
58  public final Iterable<T> children(final T root) {
59    checkNotNull(root);
60    return new FluentIterable<T>() {
61      @Override
62      public Iterator<T> iterator() {
63        return new AbstractIterator<T>() {
64          boolean doneLeft;
65          boolean doneRight;
66
67          @Override
68          protected T computeNext() {
69            if (!doneLeft) {
70              doneLeft = true;
71              Optional<T> left = leftChild(root);
72              if (left.isPresent()) {
73                return left.get();
74              }
75            }
76            if (!doneRight) {
77              doneRight = true;
78              Optional<T> right = rightChild(root);
79              if (right.isPresent()) {
80                return right.get();
81              }
82            }
83            return endOfData();
84          }
85        };
86      }
87    };
88  }
89
90  @Override
91  UnmodifiableIterator<T> preOrderIterator(T root) {
92    return new PreOrderIterator(root);
93  }
94
95  /*
96   * Optimized implementation of preOrderIterator for binary trees.
97   */
98  private final class PreOrderIterator extends UnmodifiableIterator<T>
99      implements PeekingIterator<T> {
100    private final Deque<T> stack;
101
102    PreOrderIterator(T root) {
103      this.stack = new ArrayDeque<T>();
104      stack.addLast(root);
105    }
106
107    @Override
108    public boolean hasNext() {
109      return !stack.isEmpty();
110    }
111
112    @Override
113    public T next() {
114      T result = stack.removeLast();
115      pushIfPresent(stack, rightChild(result));
116      pushIfPresent(stack, leftChild(result));
117      return result;
118    }
119
120    @Override
121    public T peek() {
122      return stack.getLast();
123    }
124  }
125
126  @Override
127  UnmodifiableIterator<T> postOrderIterator(T root) {
128    return new PostOrderIterator(root);
129  }
130
131  /*
132   * Optimized implementation of postOrderIterator for binary trees.
133   */
134  private final class PostOrderIterator extends UnmodifiableIterator<T> {
135    private final Deque<T> stack;
136    private final BitSet hasExpanded;
137
138    PostOrderIterator(T root) {
139      this.stack = new ArrayDeque<T>();
140      stack.addLast(root);
141      this.hasExpanded = new BitSet();
142    }
143
144    @Override
145    public boolean hasNext() {
146      return !stack.isEmpty();
147    }
148
149    @Override
150    public T next() {
151      while (true) {
152        T node = stack.getLast();
153        boolean expandedNode = hasExpanded.get(stack.size() - 1);
154        if (expandedNode) {
155          stack.removeLast();
156          hasExpanded.clear(stack.size());
157          return node;
158        } else {
159          hasExpanded.set(stack.size() - 1);
160          pushIfPresent(stack, rightChild(node));
161          pushIfPresent(stack, leftChild(node));
162        }
163      }
164    }
165  }
166
167  // TODO(user): see if any significant optimizations are possible for breadthFirstIterator
168
169  public final FluentIterable<T> inOrderTraversal(final T root) {
170    checkNotNull(root);
171    return new FluentIterable<T>() {
172      @Override
173      public UnmodifiableIterator<T> iterator() {
174        return new InOrderIterator(root);
175      }
176    };
177  }
178
179  private final class InOrderIterator extends AbstractIterator<T> {
180    private final Deque<T> stack;
181    private final BitSet hasExpandedLeft;
182
183    InOrderIterator(T root) {
184      this.stack = new ArrayDeque<T>();
185      this.hasExpandedLeft = new BitSet();
186      stack.addLast(root);
187    }
188
189    @Override
190    protected T computeNext() {
191      while (!stack.isEmpty()) {
192        T node = stack.getLast();
193        if (hasExpandedLeft.get(stack.size() - 1)) {
194          stack.removeLast();
195          hasExpandedLeft.clear(stack.size());
196          pushIfPresent(stack, rightChild(node));
197          return node;
198        } else {
199          hasExpandedLeft.set(stack.size() - 1);
200          pushIfPresent(stack, leftChild(node));
201        }
202      }
203      return endOfData();
204    }
205  }
206
207  private static <T> void pushIfPresent(Deque<T> stack, Optional<T> node) {
208    if (node.isPresent()) {
209      stack.addLast(node.get());
210    }
211  }
212}
213