刷题之剑指 Offer II 047. 二叉树剪枝

给定一个二叉树 根节点 root ，树的每个节点的值要么是 0，要么是 1。请剪除该二叉树中所有节点的值为 0 的子树。

节点 node 的子树为 node 本身，以及所有 node 的后代。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/pOCWxh
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

通过分析，所要满足的条件，要剪掉的子枝包括的条件是：

1. 子树根节点值为0
2. 左子树为空 / 左子树中所有节点值为0
3. 右子树为空 / 右子树中所有节点值为0

利用后序遍历，只考虑当前结点的情况：

1. 节点值为0
2. 左右子树为空

由此写出递归遍历

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        if( root == null) {
            return root;
        }
        root.left = pruneTree(root.left);
        root.right = pruneTree(root.right);
        if (root.val == 0 && root.left == null && root.right == null){
            root = null;
        }
        return root;
    }
}

如果非递归调用，则需要一个辅助栈来进行后续遍历

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        Deque<TreeNode> stack = new LinkedList<>();
        Deque<TreeNode> mark = new LinkedList<>(); //使用辅助栈
        TreeNode node = root;
        while(node != null || !stack.isEmpty()){
            while(node != null){
                stack.push(node);
                node = node.left;
            }
            while(!mark.isEmpty() && mark.peek() == stack.peek()){
                TreeNode cur = stack.pop();
                if(mark.pop().val == 0 && cur.left == null 
                && cur.right == null){
                    if(!stack.isEmpty()){
                        TreeNode tmp = stack.peek();
                        if(tmp.left == cur){
                            tmp.left = null;
                        }else{
                            tmp.right = null;
                        }
                    } else{
                        return null;
                    }
                }
            }
            if(!stack.isEmpty()){
                node = stack.peek();
                mark.push(node);
                node = node.right;
            }
        }
        return root;
    }
}