Given a binary tree, it's quite common to perform some operation on all nodes while traversing pre-order, in-order, post-order or level-order, depending on the task at hand. For example you might want to extract sorted elements from a binary search tree, for which in-order traversal is handy. Or do a breadth first search to find the shortest path to an element from the root, for which level-order traversal is handy.
Given a TreeNode defined as:
public class TreeNode<T> {
public final T value;
public TreeNode<T> left;
public TreeNode<T> right;
public TreeNode(T x) {
value = x;
}
@Override
public String toString() {
return String.valueOf(value);
}
}
The utility class to generate the various iterators:
package com.janosgyerik.examples.tree.binarytree;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Stack;
public class Iterators {
private Iterators() {
// utility class, forbidden constructor
}
public static <T> Iterator<T> preOrderIterator(TreeNode<T> root) {
return new PreOrderIterator<>(root);
}
public static <T> Iterator<T> inOrderIterator(TreeNode<T> root) {
return new InOrderIterator<>(root);
}
public static <T> Iterator<T> postOrderIterator(TreeNode<T> root) {
return new PostOrderIterator<>(root);
}
public static <T> Iterator<T> levelOrderIterator(TreeNode<T> root) {
return new LevelOrderIterator<>(root);
}
private static class PreOrderIterator<T> implements Iterator<T> {
private final Stack<TreeNode<T>> stack = new Stack<>();
private PreOrderIterator(TreeNode<T> root) {
if (root != null) {
stack.push(root);
}
}
@Override
public boolean hasNext() {
return !stack.isEmpty();
}
@Override
public T next() {
TreeNode<T> node = stack.pop();
if (node.right != null) {
stack.push(node.right);
}
if (node.left != null) {
stack.push(node.left);
}
return node.value;
}
}
private static class InOrderIterator<T> implements Iterator<T> {
private Stack<TreeNode<T>> stack = new Stack<>();
private InOrderIterator(TreeNode<T> root) {
moveToLeftMost(root);
}
private void moveToLeftMost(TreeNode<T> node) {
while (node != null) {
stack.push(node);
node = node.left;
}
}
@Override
public boolean hasNext() {
return !stack.isEmpty();
}
@Override
public T next() {
TreeNode<T> current = stack.pop();
moveToLeftMost(current.right);
return current.value;
}
}
private static class PostOrderIterator<T> implements Iterator<T> {
private Stack<TreeNode<T>> stack = new Stack<>();
private PostOrderIterator(TreeNode<T> root) {
moveToNextLeaf(root);
}
private void moveToNextLeaf(TreeNode<T> node) {
while (node != null) {
stack.push(node);
node = node.left != null ? node.left : node.right;
}
}
@Override
public boolean hasNext() {
return !stack.isEmpty();
}
@Override
public T next() {
TreeNode<T> current = stack.pop();
if (!stack.isEmpty()) {
TreeNode<T> parent = stack.peek();
if (parent.right != current) {
moveToNextLeaf(parent.right);
}
}
return current.value;
}
}
private static class LevelOrderIterator<T> implements Iterator<T> {
private final Queue<TreeNode<T>> queue = new LinkedList<>();
private LevelOrderIterator(TreeNode<T> root) {
queue.add(root);
}
@Override
public boolean hasNext() {
return !queue.isEmpty();
}
@Override
public T next() {
TreeNode<T> current = queue.poll();
if (current.left != null) {
queue.add(current.left);
}
if (current.right != null) {
queue.add(current.right);
}
return current.value;
}
}
}
Unit tests:
package com.janosgyerik.examples.tree.binarytree;
import org.junit.Test;
import java.util.Arrays;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import static org.junit.Assert.*;
public class IteratorsTest {
// example tree from http://en.wikipedia.org/wiki/Tree_traversal
/*
F
B G
A D I
C E H
*/
private final TreeNode<Character> root;
public IteratorsTest() {
root = new TreeNode<>('F');
root.left = new TreeNode<>('B');
root.left.left = new TreeNode<>('A');
root.left.right = new TreeNode<>('D');
root.left.right.left = new TreeNode<>('C');
root.left.right.right = new TreeNode<>('E');
root.right = new TreeNode<>('G');
root.right.right = new TreeNode<>('I');
root.right.right.left = new TreeNode<>('H');
}
static <T> List<T> iterateToList(Iterator<T> iterator) {
List<T> list = new LinkedList<>();
while (iterator.hasNext()) {
list.add(iterator.next());
}
return list;
}
@Test
public void testPreOrderIterator() {
Iterator<Character> iterator = Iterators.preOrderIterator(root);
assertEquals(Arrays.asList('F', 'B', 'A', 'D', 'C', 'E', 'G', 'I', 'H'), iterateToList(iterator));
}
@Test
public void testInOrderIterator() {
Iterator<Character> iterator = Iterators.inOrderIterator(root);
assertEquals(Arrays.asList('A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I'), iterateToList(iterator));
}
@Test
public void testPostOrderIterator() {
Iterator<Character> iterator = Iterators.postOrderIterator(root);
assertEquals(Arrays.asList('A', 'C', 'E', 'D', 'B', 'H', 'I', 'G', 'F'), iterateToList(iterator));
}
@Test
public void testLevelOrderIterator() {
Iterator<Character> iterator = Iterators.levelOrderIterator(root);
assertEquals(Arrays.asList('F', 'B', 'G', 'A', 'D', 'I', 'C', 'E', 'H'), iterateToList(iterator));
}
}
I'm looking for a review of all aspects of all the above code.
