Iterative tree traversal

The recursive solution is trivial, how about iterative tree traversal.

The easiest of the three (pre-order, in-order and post-order) is the pre-order traversal. The idea is simple, we use a stack to simulate the call stack of the recursive function.

public class TreeTraversal {
    public ArrayList<TreeNode> iterTraversal(TreeNode root) {
        ArrayList<TreeNode> result = new ArrayList<TreeNode>();
        Stack<TreeNode> stack = new Stack<TreeNode>();
        stack.push(root);
        while(!stack.isEmpty()){
            TreeNode curr = stack.pop();
            result.add(curr);
            //right first
            if (curr.right != null) {
                stack.push(curr.right);
            }
            if (curr.left != null) {
                stack.push(curr.left);
            }
        }
        return result;
    }
} 

Then the second easiest would be the in-order traversal, the basic idea is the same, except that the in-order traversal is like the dfs. Besides the order of putting new element into the stack, you need to have a way to keep track of those nodes that has been visited to avoid infinite loop.

public class TreeTraversal {
    public ArrayList<TreeNode> iterTraversal(TreeNode root) {
        ArrayList<TreeNode> result = new ArrayList<TreeNode>();
        Stack<TreeNode> stack = new Stack<TreeNode>();
        HashTable<Integer> table = new HashTable<Integer>(); //assume each node has a unique value
        stack.push(root);
        while(!stack.isEmpty()){
            TreeNode curr = stack.peek();
            if (curr.right != null && !table.contains(curr.right.val)) {
                stack.push(curr.right);
            }
            else if (curr.left != null && !table.contains(curr.left.val)) {
                stack.push(curr.left);
            }
            else {
                TreeNode tmp = stack.pop();
                result.add(tmp);
                table.add(tmp.val);
            }
        }
        return result;
    }
} 

However, this method will require either modify the tree to have a visited flag for each node, or require a linear space to store the visited node, which is not very desirable. Can we do this without recording which nodes are visited?

Yes. First, the current pointer is initialized to the root. Keep traversing to its left child while pushing visited nodes onto the stack. When you reach a NULL node (ie, you’ve reached a leaf node), you would pop off an element from the stack and set it to current. Now is the time to print current’s value. Then, current is set to its right child and repeat the process again. When the stack is empty, this means you’re done printing. a left subtree will be stored in the stack and will be visited in the order of the stack pop. while right subtree can be visited when curr -> curr.right

The key point here is to check if curr is null, if it is, then it means current subtree has been traversed and then you could pop another node from the stack and try traverse its right subtree.

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public List<Integer> inorderTraversal(TreeNode root) {
        List<Integer> result = new ArrayList<Integer>();
        if (root == null) 
            return result;
        Stack<TreeNode> stack = new Stack<TreeNode>();
        TreeNode curr = root;
        while(!stack.isEmpty() || curr != null) {
            if (curr != null) { // if curr has a left child, go visit left subtree
                stack.push(curr);
                curr = curr.left;
            }
            else {
                curr = stack.peek(); //when the left subtree has been traversed, curr should point to the root
                result.add(curr.val); //save the current
                stack.pop(); //the root has been visited, go to right subtree
                curr = curr.right; // if right was null ,then the current subtree has been traversed
            }
        }
        return result;
    }
}

But is there a way to do iterative in-order traversal without using a stack? Yes. use threaded binary search tree.

This is a data structure that will provide a method of constant space traversal.

For the post-order traversal, which is the post complex. If a visited flag is allowed, the solution is straight forward. However, it becomes interesting if visited flag or recording is not permitted.

The idea here is to use two references, curr and prev to determine the direction of the traversal. curr is the node that is on top of the stack and prev is the last node that has been traversed. When prev is curr‘s parent, we are traversing down the tree. In this case, we try to traverse to curr‘s left child if available (ie, push left child to the stack). If it is not available, we look at curr‘s right child. If both left and right child do not exist (ie, curr is a leaf node), we print curr‘s value and pop it off the stack.

If prev is curr‘s left child, we are traversing up the tree from the left. We look at curr‘s right child. If it is available, then traverse down the right child (ie, push right child to the stack), otherwise print curr‘s value and pop it off the stack.

If prev is curr‘s right child, we are traversing up the tree from the right. In this case, we print curr‘s value and pop it off the stack.

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public List<Integer> postorderTraversal(TreeNode root) {
        List<Integer> result = new ArrayList<Integer>();
        if (root == null)
            return result;
        Stack<TreeNode> stack = new Stack<TreeNode>();
        TreeNode curr = root;
        TreeNode prev = null;
        stack.push(root);
        while(!stack.isEmpty()) {
            //if leaf store the current into the result
            if (curr.left == null && curr.right == null) {
                TreeNode tmp = stack.pop();
                result.add(tmp.val);
                prev = curr;
                if (stack.isEmpty())
                    break;
                curr = stack.peek();
            }
            // if prev is the parent of curr, keep going down the tree
            else if (prev == null || prev.left == curr || prev.right == curr) {
                prev = curr;
                if (curr.left != null) {
                    curr = curr.left;
                }
                else if (curr.right != null) {
                    curr = curr.right;
                }
                stack.push(curr);
            }
            // if prev is the left child of curr, going to the right child
            else if (curr.left == prev) {
                
                
                if (curr.right != null) { // traverse the right subtree
                    prev = curr;
                    stack.push(curr.right);
                    curr = curr.right;    
                }
                else { // the right subtree is empty, pop the root node
                    TreeNode tmp = stack.pop();
                    result.add(tmp.val);
                    prev = curr;
                    if (stack.isEmpty())
                        break;
                    curr = stack.peek();
                }
                
            }
            // if prev is the right child of curr, pop current, going up tree
            else {
                TreeNode tmp = stack.pop();
                result.add(tmp.val);
                prev = curr;
                if (stack.isEmpty())
                    break;
                curr = stack.peek();
            }
        }
        return result;
    }
}

The other quite beautiful and magical idea to do this is to use the fact that the reverse of post-order traversal is post-order traversal, So here is how it works

An alternative solution is to use two stacks. Try to work it out on a piece of paper. I think it is quite magical and beautiful. You will think that it works magically, but in fact it is doing a reversed pre-order traversal. That is, the order of traversal is a node, then its right child followed by its left child. This yields post-order traversal in reversed order. Using a second stack, we could reverse it back to the correct order.

Here is how it works:

1)Push the root node to the first stack.
2)Pop a node from the first stack, and push it to the second stack.
3)Then push its left child followed by its right child to the first stack.
Repeat step 2) and 3) until the first stack is empty.
Once done, the second stack would have all the nodes ready to be traversed in post-order. Pop off the nodes from the second stack one by one and you’re done.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s