A General Approach to Convert Recursion to Iteration

Anything
, ,

Generally, we prefer to write recursive solutions for tree, linked list, backtracking, and DFS/BFS problems since the structure of these problems could be recursively defined.

Recursion is good choice for search, enumeration, and divide-and-conquer. — EPI

Thus, the steps that lie in the recursive function correspond to the way we think of the problem. Later, we will see how the code of preorder traversal works and then convert it to an iterative version, and finally to an optimized iterative version that we usually write if an iterative version is required.

1
2
3
4
5
6
7
// Recursive version (Easy to Read)
public void preorder(TreeNode root) {
  if (root == null) return; // base case
  foo(root); // print or ...
  preorder(root.left);
  preorder(root.right);
}

No Stack (Tail-Recursion)

In the past, many people told me that every recursive function could be rewritten as an iterative function. It is true, but when I start writing the code I usually get confused of storing parameters and intermediate return values.

As for the intermediate return values, I learned that recursion can be easily removed from a tail-recursive program by using a while-loop—no stack is needed (from EPI). Let’s take factorial computation as an example.

1
2
3
4
5
6
// Not a tail-recursive version
// returns n * (n - 1) * ... * 1
public int factorial(int n) {
  if (n == 1) return 1;
  return n * factorial(n - 1);
}

The above function is not tail-recursive because we store n in the stack frame along the way. Its value would be used after we get the returned value from factorial(n - 1).

One common way of converting a recursive function to a tail-recursive function is by passing the values we would have stored by parameters.

1
2
3
4
5
6
7
8
System.out.println(factorial(10, 1));

// A tail-recursive version
public int factorial(int n, int result) {
  if (n == 1) return result;
  result *= n;
  return factorial(n - 1, result);
}

We can also implement it iteratively without much effort.

1
2
3
4
5
6
7
8
public int factorial(int n) {
  int result = 1;
  while (n >= 1) {
    result *= n;
    n -= 1;
  }
  return result;
}

No Return Value (Tree Traversal), by Me

General Conversion

Now let’s go through the example of tree preorder traversal to illustrate how we convert a recursive function to an iterative one.

First, we need to mark states in the recursive function. We mark states at the beginning of the function and after each call of recursive calls.

1
2
3
4
5
6
7
8
9
public void preorder(TreeNode root) {
  // S1: begin
  if (root == null) return; // base case
  foo(root); // print or ...
  preorder(root.left);
  // S2: left done
  preorder(root.right);
  // S3: finish
}

Since we have one parameter, two stacks are required. The extra one is for states. In the while loop, we handle each state respectively. Here is the code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
public void preorderItr(TreeNode root) {
  // Init
  Stack<TreeNode> nodeST = new Stack<>();
  Stack<String> stateST = new Stack<>();

  while (nodeST.size() > 0) {
    // pop a node
    TreeNode x = nodeST.pop(); String s = stateST.pop();
    // handle each state respectively
    if (s.equals("S1: begin")) {
      if (x == null) { // base case
        continue; // return --> continue
      }
      foo(root); // print or ...
      nodeST.push(x); stateST.push("S2: left done");
      nodeST.push(x.left); stateST.push("S1: begin");
    } else if (s.equals("S2: left done")) {
      nodeST.push(x); stateST.push("S3: finish");
      nodeST.push(x.right); stateST.push("S1: begin");
    } else { // State 3: finish
      // at the end of the function, so do nothing
    }
  }
}

The code is quite long, but it actually looks well-organized. After “stressful” examination, you might wonder how the two lines of code below work:

1
2
nodeST.push(x); stateST.push("S2: left done");
nodeST.push(x.left); stateST.push("S1: begin");

Here is the idea. Before doing a recursive call, we should do two things:

  1. Save the current state (stateST.push(S2: left done))
  2. Push the new state (stateST.push(S1: begin))

If you think it is difficult to understand how it works, consider the fact that the state helps direct us to the specific code segment in the loop. For example, the new state always goes with the state S1: begin because we will do it from the beginning of the function.

Note: The order matters! Since stack is an FILO structure, we should push the current state such that we could retrieve it later after we handle the new state.

By now, you should have seen that the iterative version that works well. However, it looks pretty wordy to some extent. In the next section, you will see how it is converted to a version that we often encounter.

Further Optimization

The iterative version we usually encounter is as follows:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
public void preorderItr(TreeNode root) {
  // Init
  Stack<TreeNode> nodeST = new Stack<>();
  while (nodeST.size() > 0) {
    // pop a node
    TreeNode x = nodeST.pop();
    if (x == null) { // base case
      continue; // return --> continue
    }
    foo(root); // print or ...
    nodeST.push(x.right);
    nodeST.push(x.left);
  }
}

In fact, we will modify the wordy code in the previous section and get the above code.

First, we can safely remove the last state S3: finish because there is no code inside the else branch. The statement nodeST.push(x); stateST.push("S3: finish"); is also removed. Think about it why.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// ...
while (nodeST.size() > 0) {
  // pop a node
  TreeNode x = nodeST.pop(); String s = stateST.pop();
  // handle each state respectively
  if (s.equals("S1: begin")) {
    if (x == null) { // base case
      continue; // return --> continue
    }
    foo(root); // print or ...
    nodeST.push(x); stateST.push("S2: left done"); // A
    nodeST.push(x.left); stateST.push("S1: begin");
  } else { // S2: left done
    nodeST.push(x.right); stateST.push("S1: begin"); // B
  }
}
// ...

Then it comes to the most interesting part. We will delete A and put B into that spot. As mentioned before, we could safely remove the empty else if there is no code in it. Why?

Try to ponder it in this way. The intention of A is to execute the code in S2: left done. We can put the job in B over here. Notice that by doing A we actually create a signal that notifies us to do the job in B, which also creates a signal S1: begin inside the stack to notify some future execution. Therefore, we can cut the indirect process and integrate two parts. Here is the code after the modification.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
// ...
while (nodeST.size() > 0) {
  // pop a node
  TreeNode x = nodeST.pop(); String s = stateST.pop();
  // handle each state respectively
  if (s.equals("S1: begin")) {
    if (x == null) { // base case
      continue; // return --> continue
    }
    foo(root); // print or ...
    nodeST.push(x.right); stateST.push("S1: begin"); // B
    nodeST.push(x.left); stateST.push("S1: begin");
  }
}
// ...

Finally, since there is only one state, we can remove the outer if statement and also delete the stateST.

1
2
3
4
5
6
7
8
9
10
11
12
// ...
while (nodeST.size() > 0) {
  // pop a node
  TreeNode x = nodeST.pop(); String s = stateST.pop();
  if (x == null) { // base case
    continue; // return --> continue
  }
  foo(root); // print or ...
  nodeST.push(x.right); stateST.push("S1: begin"); // B
  nodeST.push(x.left); stateST.push("S1: begin");
}
// ...

By now, is there any difference between the code above and the iterative version we usually encounter?

With Return Value (Symmetric Tree), by Elon Xu

Cont’d