EaglerForge/sources/main/java/com/google/common/collect/BinaryTreeTraverser.java

/*
 * Copyright (C) 2012 The Guava Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.common.collect;

import static com.google.common.base.Preconditions.checkNotNull;

import java.util.ArrayDeque;
import java.util.BitSet;
import java.util.Deque;
import java.util.Iterator;

import com.google.common.annotations.Beta;
import com.google.common.annotations.GwtCompatible;
import com.google.common.base.Optional;

/**
 * A variant of {@link TreeTraverser} for binary trees, providing additional
 * traversals specific to binary trees.
 *
 * @author Louis Wasserman
 * @since 15.0
 */
@Beta
@GwtCompatible(emulated = true)
public abstract class BinaryTreeTraverser<T> extends TreeTraverser<T> {
	// TODO(user): make this GWT-compatible when we've checked in ArrayDeque and
	// BitSet emulation

	/**
	 * Returns the left child of the specified node, or {@link Optional#absent()} if
	 * the specified node has no left child.
	 */
	public abstract Optional<T> leftChild(T root);

	/**
	 * Returns the right child of the specified node, or {@link Optional#absent()}
	 * if the specified node has no right child.
	 */
	public abstract Optional<T> rightChild(T root);

	/**
	 * Returns the children of this node, in left-to-right order.
	 */
	@Override
	public final Iterable<T> children(final T root) {
		checkNotNull(root);
		return new FluentIterable<T>() {
			@Override
			public Iterator<T> iterator() {
				return new AbstractIterator<T>() {
					boolean doneLeft;
					boolean doneRight;

					@Override
					protected T computeNext() {
						if (!doneLeft) {
							doneLeft = true;
							Optional<T> left = leftChild(root);
							if (left.isPresent()) {
								return left.get();
							}
						}
						if (!doneRight) {
							doneRight = true;
							Optional<T> right = rightChild(root);
							if (right.isPresent()) {
								return right.get();
							}
						}
						return endOfData();
					}
				};
			}
		};
	}

	@Override
	UnmodifiableIterator<T> preOrderIterator(T root) {
		return new PreOrderIterator(root);
	}

	/*
	 * Optimized implementation of preOrderIterator for binary trees.
	 */
	private final class PreOrderIterator extends UnmodifiableIterator<T> implements PeekingIterator<T> {
		private final Deque<T> stack;

		PreOrderIterator(T root) {
			this.stack = new ArrayDeque<T>();
			stack.addLast(root);
		}

		@Override
		public boolean hasNext() {
			return !stack.isEmpty();
		}

		@Override
		public T next() {
			T result = stack.removeLast();
			pushIfPresent(stack, rightChild(result));
			pushIfPresent(stack, leftChild(result));
			return result;
		}

		@Override
		public T peek() {
			return stack.getLast();
		}
	}

	@Override
	UnmodifiableIterator<T> postOrderIterator(T root) {
		return new PostOrderIterator(root);
	}

	/*
	 * Optimized implementation of postOrderIterator for binary trees.
	 */
	private final class PostOrderIterator extends UnmodifiableIterator<T> {
		private final Deque<T> stack;
		private final BitSet hasExpanded;

		PostOrderIterator(T root) {
			this.stack = new ArrayDeque<T>();
			stack.addLast(root);
			this.hasExpanded = new BitSet();
		}

		@Override
		public boolean hasNext() {
			return !stack.isEmpty();
		}

		@Override
		public T next() {
			while (true) {
				T node = stack.getLast();
				boolean expandedNode = hasExpanded.get(stack.size() - 1);
				if (expandedNode) {
					stack.removeLast();
					hasExpanded.clear(stack.size());
					return node;
				} else {
					hasExpanded.set(stack.size() - 1);
					pushIfPresent(stack, rightChild(node));
					pushIfPresent(stack, leftChild(node));
				}
			}
		}
	}

	// TODO(user): see if any significant optimizations are possible for
	// breadthFirstIterator

	public final FluentIterable<T> inOrderTraversal(final T root) {
		checkNotNull(root);
		return new FluentIterable<T>() {
			@Override
			public UnmodifiableIterator<T> iterator() {
				return new InOrderIterator(root);
			}
		};
	}

	private final class InOrderIterator extends AbstractIterator<T> {
		private final Deque<T> stack;
		private final BitSet hasExpandedLeft;

		InOrderIterator(T root) {
			this.stack = new ArrayDeque<T>();
			this.hasExpandedLeft = new BitSet();
			stack.addLast(root);
		}

		@Override
		protected T computeNext() {
			while (!stack.isEmpty()) {
				T node = stack.getLast();
				if (hasExpandedLeft.get(stack.size() - 1)) {
					stack.removeLast();
					hasExpandedLeft.clear(stack.size());
					pushIfPresent(stack, rightChild(node));
					return node;
				} else {
					hasExpandedLeft.set(stack.size() - 1);
					pushIfPresent(stack, leftChild(node));
				}
			}
			return endOfData();
		}
	}

	private static <T> void pushIfPresent(Deque<T> stack, Optional<T> node) {
		if (node.isPresent()) {
			stack.addLast(node.get());
		}
	}
}