How can I learn to convert recursive solution to iterative solution while preserving the space complexity?

Question

I am trying to solve a easy binary tree question to convert the existing binary tree to a sum tree. The recursive solution for this is :

class Solution{
    public void toSumTree(Node root){
         solve(root);
    }
    private int solve(Node node){
        if(node == null){
            return 0;
        }
        
        int leftSubtreeValue = solve(node.left);
        int rightSubtreeValue = solve(node.right);
        int currentNodeValue = node.data;
        
        node.data = leftSubtreeValue + rightSubtreeValue;
        
        return currentNodeValue + leftSubtreeValue + rightSubtreeValue;
    }
}

This give us the time and space complexity as O(N) and O(H) respectively.

Now since the above recursive approach is similar to post-order traversal, I think we can use some version of iterative post-order travesal like below :

class Solution {
    public List<Integer> postorderTraversal(TreeNode root) { 
        List<Integer> postOrderTrvNodes = new LinkedList<>();
        if(root == null){
             return postOrderTrvNodes;
         }
        Stack<TreeNode> stack = new Stack<>();
        TreeNode current = root;
        TreeNode lastVisited = null;
        while (!stack.isEmpty() || current != null) {
            if (current != null) {
                stack.push(current);
                current = current.left;
            } else {
                TreeNode peekNode = stack.peek();
                if (peekNode.right != null && peekNode.right != lastVisited) {
                    current = peekNode.right;
                } else {
                    TreeNode poppedNode = stack.pop();
                    postOrderTrvNodes.add(poppedNode.val);
                    lastVisited = poppedNode;
                }
            }
        }
        return postOrderTrvNodes;
    }
}

What I am not able to come up with is how to store the required sum of subtree at each node so that space complexity remains O(H). If we use a HashMap that stores sum of subtree at current node in a HashMap<Node, Integer> then the space complexity will become O(N).

What I want is some protocol that can be followed in general to mimic what Java is doing with internal stack.

Answer 1

I'm not entirely sure I understand the question, but the following iterative function will return a post-order traversal of the given tree, in O(N) steps and with O(H) space.

public static List<Integer> postorderTraversal(TreeNode root) {
    final List<Integer> results = new LinkedList<>();

    final Stack<TreeNode> stack = new Stack<>();
    final Stack<Boolean> dirStack = new Stack<>();

    TreeNode curr = root;

    while (!stack.isEmpty() || curr != null) {
        if (curr == null) {
            curr = stack.pop();
            final boolean up = dirStack.pop();
            if (up) {
                results.add(curr.value);
                curr = null;
            }
        } else {
            stack.push(curr);
            dirStack.push(true);
            stack.push(curr.right);
            dirStack.push(false);
            curr = curr.left;
        }
    }

    return results;
}

A simple test:

public static void main(String[] args) {
    final TreeNode root =
            new TreeNode(
                    new TreeNode(
                            new TreeNode(1),
                            5,
                            new TreeNode(
                                    new TreeNode(2),
                                    4,
                                    new TreeNode(3))
                    ),
                    9,
                    new TreeNode(
                            new TreeNode(6),
                            8,
                            new TreeNode(7)
                    )
            );

    final List<Integer> l = postorderTraversal(root);

    System.out.println("PostorderTraversal = " + l);
}

outputs the following:

PostorderTraversal = [1, 2, 3, 4, 5, 6, 7, 8, 9]

The algorithm is derived by considering the recursive solution for a post-order tree traversal. For each node it has three steps:

1. Process left node.
2. Process right node.
3. Add node value to list.

When we convert this to an iterative solution, we need to effectively defer steps 2 & 3 by pushing them onto a stack, and then proceeed with step 1. Step 2 & 3 both require a node from the stack, but we also require a boolean flag to distinguish the two steps. I.e. we need a stack of nodes and a stack of boolean values.