【递归】二叉树剪枝

题目

814. 二叉树剪枝

---

给你二叉树的根结点 root ，此外树的每个结点的值要么是 0 ，要么是 1 。

返回移除了所有不包含 1 的子树的原二叉树。

节点 node 的子树为 node 本身加上所有 node 的后代。

示例 1：

输入：root = [1,null,0,0,1]
输出：[1,null,0,null,1]
解释：
只有红色节点满足条件“所有不包含 1 的子树”。 右图为返回的答案。

示例 2：

输入：root = [1,0,1,0,0,0,1]
输出：[1,null,1,null,1]

示例 3：

输入：root = [1,1,0,1,1,0,1,0]
输出：[1,1,0,1,1,null,1]

提示：

• 树中节点的数目在范围 [1, 200] 内
• Node.val 为 0 或 1

解题

方法一：递归

思路

这题可以用递归来做，递归返回传入节点的状态，即:是否「以该节点为根的子树」中所有节点都不含 1
返回 false 表示都不含 1，true 表示至少有一个含 1 的节点，那么对于某个节点：

• 如果该节点为空，则返回 false，此为递归的结束条件（走到空节点）。
• 否则的话，递归求出该节点左子节点的状态，和右子节点的状态。
• 如果左子节点状态为 false 删除左子树，如果右子节点状态为 false 删除右子树。
• 返回 该节点为 1 或 左子树有 1 或 右子树有 1。

最后如果根节点状态为 false 则返回 null，否则返回根节点。

代码

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        return recursive(root) ? root : null;
    }

    private boolean recursive(TreeNode node) {
        if (node == null) return false;
        boolean left = recursive(node.left);
        boolean right = recursive(node.right);
        if (!left) node.left = null;
        if (!right) node.right = null;
        return node.val == 1 || left || right;
    }
}

不借助辅助方法的递归

思路与上面大概一致。

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