Implement Iterator for BinaryTree II (Pre-order)

Implement Pre-order Iterator for Binary Tree

Suppose the data structure for a Tree node is as follows:

public class TreeNode {
  int val;
  TreeNode left;
  TreeNode right;
  TreeNode(int x) { val = x; }
}

Provide an implementation of the following interface:

public interface PreOrderBinaryTreeIterator extends Iterator<Integer> { 
  /** Returns the next integer a in the pre-order traversal of the given binary tree.
   * For example, given a binary tree below,
   *       4
   *      / \
   *     2   6
   *    / \ / \
   *   1  3 5  7
   * the outputs will be 4, 2, 1, 3, 6, 5, 7. 
   */ 
  public Integer next(); 

  /** Return true if traversal has not finished; otherwise, return false.
   */ 
  public boolean hasNext();
}

Solution

The idea is slightly different from we discussed in previous post.

In previous post, we visit left nodes on the way to get to the left-most one. But in an iterator, given a top node in the stack, it is not easy to tell whether the left node of the top one has been visited or not. We can make it work by adding flags. But there are a better way to do that.

The reason why we need to keep previous nodes is that when we complete the left sub-tree, we can go to the right sub-tree. That said, what we actually need to keep track are right nodes.

We can use Stack and handle the process more naturally. Each time when we visit a node, we push its right and left children into the stack so that we can access left subtree first and then right subtree. More specifically, we use ArrayDeque which a "resizable-array implementation of the Deque interface". It provides amortized constant time operations such as add, poll, push, and pop, etc.

 public class PreOrderBinaryTreeIteratorImpl implements PreOrderBinaryTreeIterator {  
   Stack<TreeNode> stack = new ArrayDeque<TreeNode>();  
   
   /** Constructor */  
   public PreOrderBinaryTreeIterator(TreeNode root) {  
     if (root != null) {  
       stack.push(root); // add to end of queue 
     }  
   }  
   
   /** {@inheritDoc} */  
   @Override  
   public boolean hasNext() {  
     return !stack.isEmpty();  
   }  
   
   /** {@inheritDoc} */  
   @Override  
   public Integer next() {  
     if (!hasNext()) {  
       throw new NoSuchElementException("All nodes have been visited!");  
     }  
   
     TreeNode res = stack.pop(); // retrieve and remove the head of queue 
     if (res.right != null) stack.push(res.right);  
     if (res.left != null) stack.push(res.left);  
   
     return res.val;  
   }  
   
   @Override  
   public void remove() {  
     throw new UnsupportedOperationException("remove() is not supported.");  
   }  
 }  

This iterator takes extra spaces for the stack, which is O(h) at worst case, where h is the height of the tree.

With this iterator in hand, an pre-order traversal of a binary tree can be implemented as follows.

 public ArrayList<Integer> preorderTraversal(TreeNode root) {  
   PreOrderBinaryTreeIterator iterator = new PreOrderBinaryTreeIteratorImpl(root);  
   ArrayList<Integer> results = new ArrayList<Integer>();  
   while (iterator.hasNext()) {  
     results.add(iterator.next());  
   }  
   return results;  
 }