刷题升级之路：Leetcode——814.二叉树剪枝

题目描述

给你二叉树的根结点 root ，此外树的每个结点的值要么是 0 ，要么是 1 。
返回移除了所有不包含 1 的子树的原二叉树。
节点 node 的子树为 node 本身加上所有 node 的后代。

题目中给出了二叉树结点的定义：

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */

算法思路与分析

采用递归算法，我一开始的想法是，写一个函数处理以root为根结点的子树：需要判断其左子树和右子树是否有不为0的结点，全为0则删去，最后返回值加上根结点的值。递归此函数可以得出最后的结果。（特殊情况：对于结点全部为0的树，做判断后返回空指针即可）

代码实现

class Solution {
public:
    bool operate(TreeNode* root){
        int numleft = 0, numright = 0;
        if(root->left){
            numleft = operate(root->left);
        }
        if(root->right){
            numright = operate(root->right);
        }
        if(numleft==0){
            root->left = nullptr;
        }
        if(numright==0){
            root->right = nullptr;
        }
        return numleft | numright | root->val;
    }
    TreeNode* pruneTree(TreeNode* root) {
        int ans = operate(root);
        if(ans==0)
            return nullptr;
        return root;
    }
};

后来我发现直接对pruneTree函数做递归就可以得出结果：

    TreeNode* pruneTree(TreeNode* root){    // 处理以root为根结点的树
        if(!root)
            return nullptr;
        root->left = pruneTree(root->left);
        root->right = pruneTree(root->right);
        if(!root->left&& !root->right&& !root->val)
            root = nullptr;
        return root;
    }