描述

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

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

示例

示例1：

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

1028_2.png

示例2：

输入: [1,0,1,0,0,0,1]
输出: [1,null,1,null,1]
解释:

1028_1.png
示例3：

输入: [1,1,0,1,1,0,1,0]
输出: [1,1,0,1,1,null,1]
解释:

1028.png

提示

  • 二叉树的节点个数的范围是 [1,200]
  • 二叉树节点的值只会是 0 或 1

解题思路

递归：

我们可以使用递归来解决这个问题。我们用 containsOne(node) 函数来判断以 node 为根的子树中是否包含 1，其不包含 1 当且仅当以 node 的左右孩子为根的子树均不包含 1，并且 node 节点本身的值也不为 1。

如果 node 的左右孩子为根的子树不包含 1，那我们就需要把对应的指针置为空。例如当 node 的左孩子为根的子树不包含 1 时，我们将 node.left 置为 null。

在递归结束之后，如果整颗二叉树都不包含 1，那么我们返回 null，否则我们返回原来的根节点。

代码

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

    public boolean containsOne(TreeNode node) {
        if (node == null) return false;
        boolean a1 = containsOne(node.left);
        boolean a2 = containsOne(node.right);
        if (!a1) node.left = null;
        if (!a2) node.right = null;
        return node.val == 1 || a1 || a2;
    }
}

复杂度分析

时间复杂度：O(N)，其中 N 是树中节点的个数。
空间复杂度：O(H)，其中 H 是树的高度，为我们在递归时使用的栈空间大小。