牛客/LeetCode——二叉树的剪枝

二叉树的剪枝
给定二叉树根结点 root ，此外树的每个结点的值要么是 0，要么是 1。
返回移除了所有不包含 1 的子树的原二叉树。

class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        if(root==NULL)
            return root;
        _pruneTree(root);
        return root;
    }
    void _pruneTree(TreeNode* root) {
        if(root==NULL)
            return;
        //如果它的左子树满足全为0——即没有含1的节点，将其左孩子置为空
        if(_ZeroTree(root->left))
            root->left=NULL;
        //如果它的右子树满足全为0——将其右孩子置为空
        if(_ZeroTree(root->right))
            root->right=NULL;
        _pruneTree(root->left);
        _pruneTree(root->right);
    }

    //全为0的子树
    bool _ZeroTree(TreeNode* root) 
    {
        if(root==NULL)
            return true;
        if(root->val==1)
            return false;
        else
            return _ZeroTree(root->left)&&_ZeroTree(root->right);
    }
};