二叉樹遍歷的迭代寫法

阿新 • • 發佈：2021-08-28

二叉樹遍歷的迭代寫法

我們知道二叉樹的遞迴寫法非常的簡單，因為遞迴的時候隱式的維護了一個棧。而在迭代的時候也可以採用遞迴的思想，完全使用一個棧手動模擬遞迴的過程。
在其中應該特別注意節點進棧的順序和我們遍歷二叉樹所需要的順序之間的關係。

前序遍歷

在前序遍歷中，二叉樹節點的遍歷順序是父 -> 左 -> 右，這樣要注意比如儲存左節點的時候要注意該節點的父節點一定要先儲存了。
因為方法中進棧的順序是父 -> 左 -> 右，所以前序遍歷儲存結果的時候直接在進棧的時候儲存即可。

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public List<Integer> preorderTraversal(TreeNode root) {
        List<Integer> res = new ArrayList<>();
        Stack<TreeNode> stack = new Stack<>();
        
        // 完全按照遞迴的寫法來改，使用迴圈代替函式呼叫，用棧來模擬函式呼叫棧保護現場和恢復現場
        while(root != null || !stack.empty()){
            // 類似遞迴的判定條件和遍歷左節點
            while(root != null) {
                res.add(root.val);
                stack.push(root);
                root = root.left;
            }

            // 類似該節點左子樹遍歷完，返回後需要出棧恢復現場
            //（其實真正要等右子樹遍歷完才出棧，因為右子樹遍歷完之後不需要使用父節點了，前序遍歷父節點在右結點前已經儲存了，所以不影響答案）
            root = stack.peek();
            stack.pop();
            // 類似遍歷右節點，因為在遞迴中遍歷右節點還是使用遞迴函式，所以直接迴圈開始即可。
            root = root.right;
        }

        return res;
    }
}

中序遍歷

類似於前序遍歷，不過儲存的順序為左 -> 父 -> 右

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public List<Integer> inorderTraversal(TreeNode root) {
        List<Integer> res = new ArrayList<>();
        Stack<TreeNode> stack = new Stack<>();

        // 離開就進棧保護現場，回來就出棧恢復現場
        while(root != null || !stack.empty()){
            // 如果root不等於null，我就可以一直走左邊的結點
            while(root != null){
                // 因為要到另一個結點了，使用棧保護現場
                stack.push(root);
                root = root.left;
            }

            // 出棧恢復現場
            root = stack.peek();
            stack.pop();
            res.add(root.val);

            // 做右節點
            root = root.right;
        }

        return res;
    }
}

後序遍歷

在後續遍歷的過程中，因為我們儲存順序是左右父，所以要注意在第一次返回到父節點的時候，因為我們右節點先於父節點儲存，所以如果父節點的右子樹還未遍歷，那麼就不應該直接出棧，如果直接出棧，那麼父節點就丟掉了。所以我們通過判定 root.right == null || root.right == prev（prev表示上一個儲存的結果是什麼，如果該節點的右節點已經儲存了，那麼在後序遍歷中該節點的右子樹都已經儲存了，那麼就可以出棧並儲存父節點了）的話證明右子樹要麼不存在不用遍歷，要麼已經遍歷完成了。那麼此時父節點才可以真正的出棧，並儲存。

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public List<Integer> postorderTraversal(TreeNode root) {
        List<Integer> res = new ArrayList<>();
        Stack<TreeNode> stack = new Stack<>();
        TreeNode prev = null;
        while (root != null || !stack.isEmpty()) {
            while(root != null){
                stack.push(root);
                root = root.left;
            }

            root = stack.pop();
            if(root.right == null || root.right == prev){
                res.add(root.val);
                prev = root;
                root = null;
            }
            else{
                stack.push(root);
                root = root.right;
            }
        }
        
        return res;
    }
}

總結：在將遞迴程式碼改寫成迭代程式碼的時候一定要注意進棧順序和遍歷儲存順序之間關聯，和儲存時注意時候的順序。

如有錯誤，歡迎指正！

二叉樹遍歷的迭代寫法