Binary Tree Pruning

Given the root of a binary tree, return the same tree where every subtree (of the given tree) not containing a 1 has been removed. A subtree of a node node is node plus every node that is a descendant of node. For example:

Example 1: 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.

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

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

Constraints: The number of nodes in the tree is in the range [1, 200]. Node.val is either 0 or 1.

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def pruneTree(root: TreeNode) -> TreeNode:
    def contains_one(node: TreeNode) -> bool:
        if not node:
            return False

        left_contains_one = contains_one(node.left)
        right_contains_one = contains_one(node.right)

        if not left_contains_one:
            node.left = None
        if not right_contains_one:
            node.right = None

        return node.val == 1 or left_contains_one or right_contains_one

    if not contains_one(root):
        return None
    return root

Explanation:

1. TreeNode Class: Defines the structure for a binary tree node.
2. pruneTree(root: TreeNode) -> TreeNode: Main function to prune the tree.
   • It uses a helper function contains_one(node: TreeNode) -> bool to recursively check if a subtree contains at least one node with value 1.
   • If the entire tree does not contain 1, it returns None.
   • Otherwise, it returns the pruned root.
3. contains_one(node: TreeNode) -> bool: Recursive helper function.
   • Base Case: If the node is None, return False.
   • Recursively calls itself for the left and right subtrees.
   • If the left subtree does not contain 1, set node.left = None.
   • If the right subtree does not contain 1, set node.right = None.
   • Returns True if the current node's value is 1 or if either the left or right subtrees contain 1, otherwise returns False.

Example:

Consider the input root = [1,null,0,0,1]

1. The pruneTree function is called with the root node (value 1).
2. The contains_one function is called on the root:
   • It recursively calls contains_one on the left child (None), which returns False.
   • It recursively calls contains_one on the right child (value 0).
     • The right child's contains_one calls contains_one on its left child (value 0) which returns False and sets node.left = None
     • The right child's contains_one calls contains_one on its right child (value 1) which returns True
     • Since right child has a 1 in its subtree, it returns True.
   • Since the right subtree contains 1, nothing is changed.
   • The root returns True because its own value is 1.
3. The pruneTree function returns the modified root.

Big O Analysis:

• Time Complexity: O(N), where N is the number of nodes in the tree. This is because the contains_one function visits each node once.
• Space Complexity: O(H), where H is the height of the tree. This space is used by the recursion stack. In the worst case (skewed tree), H can be equal to N, resulting in O(N) space complexity. In the best case (balanced tree), H is log(N), resulting in O(log(N)) space complexity.

Edge Cases:

1. Empty Tree: If the input root is None, the function should return None.
2. Tree with Only Zeroes: If the tree contains only nodes with value 0, the function should return None.
3. Single Node Tree: If the tree contains only one node, the function should return the node if its value is 1, otherwise None if its value is 0.
4. Skewed Tree: A tree where all nodes are either left children or right children. The algorithm should handle this case without stack overflow issues (within reasonable limits).
5. Balanced Tree: A tree where the height of left and right subtrees of every node differ by at most 1.