二叉樹非遞歸遍歷的通用解法

最近在跟慕課網上《玩轉算法面試》這門課，老師講的很不錯，其中有一節提到了用棧來模擬二叉樹的前序，中序，後序遍歷。
二叉樹的遍歷的遞歸解法很簡單，但有時面試或做題時往往會要求寫出非遞歸的形式，一般教科書上實現的前序，中序和後序遍歷的非遞歸形式差別都比較大，特別是後序遍歷，很難理解。那有沒有一種通用的迭代解法能適於這三種遍歷呢？下面介紹一下這種通用解法。

基本思想

先給出前序遍歷的遞歸形式代碼如下：

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類來表示一條命令，其代碼如下：

private class Command {
	String s;    // go, print
    TreeNode node;
    Command(String s, TreeNode node) {
    	this.s = s;
        this.node = node;
    }
}

其中print代表打印結點的命令，go代表遞歸調用。
整個前序遍歷非遞歸的代碼如下：

    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));
            }
        }
	}

寫出了前序遍歷的非遞歸代碼後，中序和後序遍歷就比較簡單了，只需要交換一下入棧的順序即可。

中序非遞歸代碼

	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));放到了中間。

後序非遞歸代碼

	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題去實踐一下。

參考：慕課網《玩轉算法面試》