814. 二叉树剪枝#

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

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

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

7

8
示例 1：
9

10

11
输入：root = [1,null,0,0,1]
12
输出：[1,null,0,null,1]
13
解释：
14
只有红色节点满足条件"所有不包含 1 的子树"。 右图为返回的答案。
15
示例 2：
16

17

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

21
示例 3：
22

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

26
提示：
27

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

1
/**
2
 * Definition for a binary tree node.
3
 * struct TreeNode {
4
 *     int val;
5
 *     TreeNode *left;
6
 *     TreeNode *right;
7
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
8
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
9
 *     TreeNode(int x, TreeNode *left, TreeNode *right) :
10
        val(x), left(left), right(right) {}
11
 * };
12
 */
13
class Solution {
14
public:
15
    int dfs(TreeNode* root) {
16
        if(!root) return 0;
17
        auto l = dfs(root -> left);
18
        auto r = dfs(root -> right);
19
        if(!l) root -> left = nullptr;
20
        if(!r) root -> right = nullptr;
21
        return root -> val + l + r;
22
    }
23
    TreeNode* pruneTree(TreeNode* root) {
24
        return dfs(root) != 0 ? root : nullptr;
25
    }
26
};