814. 二叉树剪枝
题目描述
输入二叉树的根节点 root，每个节点的值只会是 0 或 1。树用层序数组表示，空位记为 null。
一个节点的子树包括这个节点本身和它的全部后代。把所有不含值为 1 的节点的子树从树上删掉。
返回剪完之后的树的根。如果整棵树都不含 1，返回空树。
示例 1
输入：root = [1,null,0,0,1] 输出：[1,null,0,null,1]
解释：根的右孩子 0 下面挂着的左孩子 0 不含 1，被剪掉；右孩子 0 因为下面有 1 而保留。
示例 2
输入：root = [1,0,1,0,0,0,1] 输出：[1,null,1,null,1]
解释：所有只含 0 的子树都被剪掉，只剩一条全是 1 的右链。
示例 3
输入：root = [1,1,0,1,1,0,1,0] 输出：[1,1,0,1,1,null,1]
解释：两个值为 0 的叶子被剪掉；根的右孩子 0 因为下面有 1 而保留。
数据范围
- 节点数在
[1, 200]范围内。 Node.val只会是0或1。
解题思路
一个子树要不要剪，取决于它的子树剪完之后还剩什么，所以用后序遍历：先递归剪左右子树，再看当前节点。若当前值为 0 且左右都已被剪空，说明整个子树不含 1，返回 None 让父节点把它摘掉；否则保留。
为什么正确
pruneTree(x) 返回后，以 x 为根的子树中每个剩余子树都至少含一个 1。
复杂度分析
时间 O(n)，每个节点访问一次；空间 O(h)，h 为树高（递归栈）。
易错点
先判断当前节点再递归（先序）是错的：一个 0 节点在剪掉孩子之前看起来有孩子，剪完后却可能变成要删的叶子。必须先剪子树。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。
class Solution:
def pruneTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
if root is None:
return None
root.left = self.pruneTree(root.left)
root.right = self.pruneTree(root.right)
if root.val == 0 and root.left is None and root.right is None:
return None
return rootclass Solution {
public TreeNode pruneTree(TreeNode root) {
if (root == null) {
return null;
}
root.left = pruneTree(root.left);
root.right = pruneTree(root.right);
if (root.val == 0 && root.left == null && root.right == null) {
return null;
}
return root;
}
}class Solution {
public:
TreeNode* pruneTree(TreeNode* root) {
if (root == nullptr) {
return nullptr;
}
root->left = pruneTree(root->left);
root->right = pruneTree(root->right);
if (root->val == 0 && root->left == nullptr && root->right == nullptr) {
return nullptr;
}
return root;
}
};func pruneTree(root *TreeNode) *TreeNode {
if root == nil {
return nil
}
root.Left = pruneTree(root.Left)
root.Right = pruneTree(root.Right)
if root.Val == 0 && root.Left == nil && root.Right == nil {
return nil
}
return root
}var pruneTree = function(root) {
if (root === null) {
return null;
}
root.left = pruneTree(root.left);
root.right = pruneTree(root.right);
if (root.val === 0 && root.left === null && root.right === null) {
return null;
}
return root;
};function pruneTree(root: TreeNode | null): TreeNode | null {
if (root === null) {
return null;
}
root.left = pruneTree(root.left);
root.right = pruneTree(root.right);
if (root.val === 0 && root.left === null && root.right === null) {
return null;
}
return root;
}