Binary Tree Pruning Python Solution with Tests

lcpy gen -n 814   # by problem number
lcpy gen -s binary_tree_pruning   # by problem name

Input: root = [1,null,0,0,1]
Output: [1,null,0,null,1]
Explanation:
Only the red nodes satisfy the property "every subtree not containing a 1".
The diagram on the right represents the answer.

Input: root = [1,0,1,0,0,0,1]
Output: [1,null,1,null,1]

Input: root = [1,1,0,1,1,0,1,0]
Output: [1,1,0,1,1,null,1]

from leetcode_py import TreeNode


class Solution:
    # Time: O(n)
    # Space: O(h)
    def prune_tree(self, root: TreeNode[int] | None) -> TreeNode[int] | None:
        if root is None:
            return None
        root.left = self.prune_tree(root.left)
        root.right = self.prune_tree(root.right)
        if root.val == 0 and root.left is None and root.right is None:
            return None
        return root