binary-tree-pruning.py

binary-tree-pruning.py

class Solution:
    def pruneTree(self, root: TreeNode) -> TreeNode:
        if root == None:
            return None
        left = self.pruneTree(root.left)
        right = self.pruneTree(root.right)
        if root.val == 1 or left or right: # not prunable
            root = TreeNode(root.val)
            root.left = left
            root.right = right
            return root
        else:
            return None