二叉树遍历的迭代写法

二叉树遍历的迭代写法

我们知道二叉树的递归写法非常的简单，因为递归的时候隐式的维护了一个栈。而在迭代的时候也可以采用递归的思想，完全使用一个栈手动模拟递归的过程。
在其中应该特别注意节点进栈的顺序和我们遍历二叉树所需要的顺序之间的关系。

前序遍历

在前序遍历中，二叉树节点的遍历顺序是父 -> 左 -> 右，这样要注意比如保存左节点的时候要注意该节点的父节点一定要先保存了。
因为方法中进栈的顺序是父 -> 左 -> 右，所以前序遍历保存结果的时候直接在进栈的时候保存即可。

/**
 * 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;
    }
}

总结：在将递归代码改写成迭代代码的时候一定要注意进栈顺序和遍历保存顺序之间关联，和保存时注意时候的顺序。