leetcode_814. 二叉树剪枝（C++，双百解法）

题目

原题链接

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

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

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

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

示例2:
输入: [1,0,1,0,0,0,1]
输出: [1,null,1,null,1]

示例3:
输入: [1,1,0,1,1,0,1,0]
输出: [1,1,0,1,1,null,1]

说明:

• 给定的二叉树最多有 100 个节点。
• 每个节点的值只会为 0 或 1 。

思路

这题的关键就在于：

怎么确认一棵子树是否都为 0 呢？

一般人的想法是——递归！如果递归的返回值都为零的话，那整棵树就都是 0 了。

这是一个很常规的思路，也不是不行，但让我们思考一下，有没有更简洁的思路呢？其实是有的，经过思考，我们可以发现全部为 0 的树的本质，也就是“所有节点之和为零”。

利用这一点，就可以很容易地写出下面的代码：

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        if(!root) return nullptr;
        if(sum(root) == 0) return nullptr; //核心代码
        TreeNode* l = pruneTree(root -> left); //保存左子树处理结果
        TreeNode* r = pruneTree(root -> right); //保存右子树处理结果
        root -> left = l; //修改左子树
        root -> right = r; //修改右子树
        return root; //返回修剪后的树
    }
    int sum(TreeNode* root) {
        if(!root) return 0; //空树之和为零
        int l = sum(root -> left), 
            r = sum(root -> right); //递归计算左右子树的节点总和
        return l + r + root -> val;
    }
};