Leetcode 814. Binary Tree Pruning

dfs

要点是这一句:

return node.val==1 or node.left or node.right

完整代码:

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None
import functools
class Solution:
    @functools.lru_cache()
    def pruneTree(self, root: TreeNode) -> TreeNode:
        def dfs(node:TreeNode)->bool:
            if not node:
                return False
            if not dfs(node.left):
                node.left=None
            if not dfs(node.right):
                node.right=None
            return node.val==1 or node.left or node.right
        if not root:
            return
        return root if dfs(root) else None