题目难度: 中等
今天继续更新 Leetcode 的剑指 Offer（专项突击版）系列, 大家在公众号 算法精选 里回复
剑指offer2就能看到该系列当前连载的所有文章了, 记得关注哦~
题目描述
给定一个二叉树 根节点 root ，树的每个节点的值要么是 0，要么是 1。请剪除该二叉树中所有节点的值为 0 的子树。
节点 node 的子树为 node 本身，以及所有 node 的后代。
示例 1:
-
输入: [1,null,0,0,1]
-
输出: [1,null,0,null,1]
-
解释: 只有红色节点满足条件“所有不包含 1 的子树”。 右图为返回的答案。
示例 2:
- 输入: [1,0,1,0,0,0,1]
- 输出: [1,null,1,null,1]
- 解释:
示例 3:
- 输入: [1,1,0,1,1,0,1,0]
- 输出: [1,1,0,1,1,null,1]
- 解释:
提示:
- 二叉树的节点个数的范围是 [1,200]
- 二叉树节点的值只会是 0 或 1
题目思考
- 如何在得到子树所有节点值的同时, 进行剪除操作?
解决方案
思路
- 分析题目, 如果某个子树的所有节点都是 0, 则其根节点的值为 0, 且左右子树的所有节点也都是 0
- 我们可以根据这一点, 使用 DFS 来判断, 具体做法如下:
- dfs 函数传入节点 node 作为参数, 返回以该节点为根的子树的所有节点是否全为 0
- 如果传入是空节点, 则返回 true, 作为递归出口
- 否则先递归调用 dfs 函数, 传入其左右子节点, 从而得到左右子树是否全为 0 的两个布尔值: leftAllZero 和 rightAllZero
- 那么当前子树要想全为 0, 就需要
node.val==0 && leftAllZero && rightAllZero, 返回它即可
- 上述步骤只是得出了每个子树的所有节点是否全 0, 题目还要求剪除全 0 子树, 如何操作呢?
- 我们可以在上面的第三步和第四步之间加入剪除操作:
- 如果左子树全 0, 就断开当前节点和其左子节点的连接
- 如果右子树全 0, 就断开当前节点和其右子节点的连接
- 这样就在得到子树所有节点值的同时, 进行了剪除操作
- 需要注意有可能整棵树的所有节点全为 0, 此时剪除后需要返回空节点
- 我们可以引入一个哨兵节点, 将其左子节点指向整棵树的根节点
- 然后从哨兵节点开始判断, 最终返回哨兵的左子节点, 这样就无需针对根节点做特殊处理了
- 下面的代码就对应了上面的整个过程, 并且有详细的注释, 方便大家理解
复杂度
- 时间复杂度 O(N): 需要遍历每个节点一遍
- 空间复杂度 O(H): H 是树的高度, 也是递归栈的空间消耗
代码
class Solution:
def pruneTree(self, root: TreeNode) -> TreeNode:
# 使用哨兵节点, 将其左子节点指向root, 并从哨兵开始判断
# 因为有可能整个树都是0, 此时需要返回空节点, 如果不用哨兵从root开始判断, 就需要额外的逻辑特殊处理root
dummy = TreeNode(0, root)
def allZero(node):
if not node:
# 递归出口, 空节点的值不是1, 所以返回True
return True
leftAllZero = allZero(node.left)
rightAllZero = allZero(node.right)
if leftAllZero:
# 左子树全为0, 断开当前节点与其的连接
node.left = None
if rightAllZero:
# 右子树全为0, 断开当前节点与其的连接
node.right = None
# 当前子树全为0需要满足: 当前节点值为0 && 左子树全为0 && 右子树全为0
return node.val == 0 and leftAllZero and rightAllZero
allZero(dummy)
# 最终哨兵节点的左子节点即为删除全0子树后的根节点 (可能为空)
return dummy.left
大家可以在下面这些地方找到我~😊
我的公众号: 算法精选, 欢迎大家扫码关注~😊
京公网安备 11010502036488号