跳至主要內容

剑指 Offer II 047. 二叉树剪枝

ITCharge小于 1 分钟

剑指 Offer II 047. 二叉树剪枝open in new window

  • 标签:树、深度优先搜索、二叉树
  • 难度:中等

题目链接

题目大意

给定一棵二叉树的根节点 root,树的每个节点值要么是 0,要么是 1

要求:剪除该二叉树中所有节点值为 0 的子树。

  • 节点 node 的子树为: node 本身,以及所有 node 的后代。

解题思路

定义辅助方法 containsOnlyZero(root) 递归判断以 root 为根的子树中是否只包含 0。如果子树中只包含 0,则返回 True。如果子树中含有 1,则返回 False。当 root 为空时,也返回 True

然后递归遍历二叉树,判断当前节点 root 是否只包含 0。如果只包含 0,则将其置空,返回 None。否则递归遍历左右子树,并设置对应的左右指针。

最后返回根节点 root

代码

class Solution:
    def containsOnlyZero(self, root: TreeNode):
        if not root:
            return True
        if root.val == 1:
            return False
        return self.containsOnlyZero(root.left) and self.containsOnlyZero(root.right)

    def pruneTree(self, root: TreeNode) -> TreeNode:
        if not root:
            return root
        if self.containsOnlyZero(root):
            return None

        root.left = self.pruneTree(root.left)
        root.right = self.pruneTree(root.right)
        return root