最近在跟慕课网上《玩转算法面试》这门课，老师讲的很不错，其中有一节提到了用栈来模拟二叉树的前序，中序，后序遍历。
二叉树的遍历的递归解法很简单，但有时面试或做题时往往会要求写出非递归的形式，一般教科书上实现的前序，中序和后序遍历的非递归形式差别都比较大，特别是后序遍历，很难理解。那有没有一种通用的迭代解法能适于这三种遍历呢？下面介绍一下这种通用解法。

基本思想

先给出前序遍历的递归形式代码如下：

1
2
3
4
5
6
7
public void preorder(TreeNode node) {
	if(node != null) {
		System.out.println(node.val);
		preorder(node.left);
		preorder(node.right);
	}
}

假设有这样一颗二叉树：

上述递归代码中其实只包含三条指令：打印结点，访问左孩子，访问右孩子，每条指令用一种命令来表示：

由于栈具有后进先出的特性，所以上述递归代码等价于向栈中推入三条命令：

上述是遍历结点1的情况，首先打印结点即count 1出栈，然后go 1-L出栈递归访问其左孩子结点2，对于结点2，也会有如上三条命令，将其也按照待访问顺序加入栈中：

依次执行上述过程直到栈为空为止。

前序遍历非递归代码

根据上述过程，我们将其写成代码，可以用一个Command类来表示一条命令，其代码如下：

1
2
3
4
5
6
7
8
private class Command {
	String s;    // go, print
    TreeNode node;
    Command(String s, TreeNode node) {
    	this.s = s;
        this.node = node;
    }
}

其中print代表打印结点的命令，go代表递归调用。
整个前序遍历非递归的代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
   private class TreeNode {
       int val;
       TreeNode left;
       TreeNode right;
       TreeNode(int x) { val = x; }
   }

   private class Command {
       String s;    // go, print
       TreeNode node;
       Command(String s, TreeNode node) {
           this.s = s;
           this.node = node;
       }
   }
public void preorderTraversal(TreeNode root) {
 	Stack<Command> st = new Stack<>();
       st.add(new Command("go", root));

       while(!st.isEmpty()) {
           Command cd = st.pop();
           if(cd.s.equals("print")) {
               System.out.println(cd.node.val)
           } else {
               assert cd.s.equals("go");
               if(cd.node.right != null) {
                   st.add(new Command("go", cd.node.right));
               }
               if(cd.node.left != null) {
                   st.add(new Command("go", cd.node.left));
               }
               st.push(new Command("print", cd.node));
           }
       }
}

写出了前序遍历的非递归代码后，中序和后序遍历就比较简单了，只需要交换一下入栈的顺序即可。

中序非递归代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
public void inorderTraversal(TreeNode root) {
 	Stack<Command> st = new Stack<>();
       st.add(new Command("go", root));

       while(!st.isEmpty()) {
           Command cd = st.pop();
           if(cd.s.equals("print")) {
               System.out.println(cd.node.val)
           } else {
               assert cd.s.equals("go");
               if(cd.node.right != null) {
                   st.add(new Command("go", cd.node.right));
               }
               st.push(new Command("print", cd.node));
               if(cd.node.left != null) {
                   st.add(new Command("go", cd.node.left));
               }
           }
       }
}

可以看到只是将st.push(new Command("print", cd.node));放到了中间。

后序非递归代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
public void postorderTraversal(TreeNode root) {
 	Stack<Command> st = new Stack<>();
       st.add(new Command("go", root));

       while(!st.isEmpty()) {
           Command cd = st.pop();
           if(cd.s.equals("print")) {
               System.out.println(cd.node.val)
           } else {
               assert cd.s.equals("go");
               st.push(new Command("print", cd.node));
               if(cd.node.right != null) {
                   st.add(new Command("go", cd.node.right));
               }
               if(cd.node.left != null) {
                   st.add(new Command("go", cd.node.left));
               }
           }
       }
}

前序，中序和后序非递归遍历的整个形式保持一致，非常好记。可以通过LeetCode上的144和145题去实践一下。

参考：慕课网《玩转算法面试》