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

文章目录

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;}
};