【二叉树的深搜】二叉树剪枝

文章目录

814. 二叉树剪枝
解题思路：深度优先遍历 + 后序遍历
另一种写法

814. 二叉树剪枝

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

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

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

示例 1：

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

示例 2：

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

示例 3：

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

提示：

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

解题思路：深度优先遍历 + 后序遍历

首先要理解这道题的意思，不是说当前节点为零就要剪枝，而是 要看以当前节点为根节点的子树是否都为零，如果是的话才进行剪枝，这个是这道题容易搞错的点！

那么既然我们需要知道当前节点的左右子树情况才进行剪枝，那么势必就要使用 后序遍历 才行，因为后序遍历是最后才处理当前节点，此时左右子树是已经处理完了的，那么我们可以让左右子树直接返回一个指针，以左子树为例，如果左子树整棵子树都是零的话，那么直接返回空指针即可，对于右子树也是如此！

此时后序遍历之后我们也拿到了左右子树返回的两个指针，我们判断一下，如果当前节点为零，并且左右子树都是空指针的话，说明当前节点也是需要被剪枝的，则直接返回一个 nullptr 即可，否则的话返回当前节点给上一层，以此类推！

所以下面分三步来进行讨论：

函数头的设计：
- 因为我们需要左右子树和当前节点返回一个节点指针，所以返回值就是 TreeNode*。然后因为整个过程其实我们只需要使用到当前节点，所以只需要一个变量来表示当前节点，那么题目给的函数头刚好符合要求，我们就直接拿题目的函数头进行使用，如下所示：
```
TreeNode* pruneTree(TreeNode* root);
```
函数体的内容：
- 首先毋庸置疑就是要进行后序遍历，先递归到左子树，拿到其返回值，判断一下返回值是否为空，为空的话表示左子树就是需要剪枝的，则直接将 root->left 置为空即可，对于右子树来说也是如此。
- 然后处理完左右子树之后，此时我们 判断一下当前节点是否为零，是的话 再判断一下左右子树是否都为空：
  - 如果都为空表示当前节点是需要剪枝的，直接返回一个 nullptr 即可！
  - 如果左右子树至少有一个不为空的话，则说明以当前节点为根节点的子树是不需要剪枝的，则返回当前节点给上一层即可！
递归函数出口：
- 出口很简单，就是递归到空节点了，直接返回一个 nullptr 即可！

/**
 * 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) {}
 * };
 */
class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        // 递归函数出口
        if(root == nullptr)
            return nullptr;

        // 后序遍历，先处理左右子树，才能知道当前节点是否需要剪枝
        root->left = pruneTree(root->left);
        root->right = pruneTree(root->right);

        // 再处理当前节点，如果当前节点为零并且是叶子节点的话，则返回false即可
        if(root->val == 0 && root->left == nullptr && root->right == nullptr)
        {
            delete root;
            root = nullptr;
        }
        return root;
    }
};

另一种写法

当然，我们也可以把 dfs() 函数单独拿出来，然后用返回值为布尔值的形式进行处理，大体思路都是一样的，只不过函数头不同而导致细节变了一点而已！

/**
 * 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) {}
 * };
 */
class Solution {
public:
    TreeNode* pruneTree(TreeNode* root) {
        if(dfs(root))
            return root;
        return nullptr;
    }

    bool dfs(TreeNode* root)
    {
        if(root == nullptr)
            return false;
        
        // 后序遍历，先处理左右子树，才能知道当前节点是否需要剪枝
        // 如果左右孩子递归处理之后发现可以剪枝的，直接将其置为nullptr即可
        if(dfs(root->left) == false)
            root->left = nullptr;
        if(dfs(root->right) == false)
            root->right = nullptr;
        
        // 再处理当前节点，如果当前节点为零并且是叶子节点的话，则返回false即可
        if(root->val == 0 && root->left == nullptr && root->right == nullptr)
        {
            delete root; // 别忘了释放一下节点防止内存泄漏
            return false;
        }
        return true;
    }
};