algoadvance

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.

Example:

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 left represents the given tree [1,null,0,0,1], where "null" stands for a null node.

     1                         1
      \                         \
       0      -->     becomes    0
      / \                         \
     0   1                         1

Constraints:

• The number of nodes in the tree is in the range [1, 200].
• The value of each node is either 0 or 1.

Clarifying Questions

1. Are there cases where the input tree might be empty?
   • No, as per the constraints, there will be at least one node.
2. What should be the format of the output?
   • The output should be the root of the pruned binary tree.

Code

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
    
    return root if contains_one(root) else None

Strategy

1. Recursive Approach:
   • Use a helper function contains_one to determine if a subtree contains a 1.
   • For each node, recursively check its left and right children.
   • If a subtree (left or right) does not contain a 1, prune (remove) that subtree by setting the corresponding child to None.
   • The function contains_one returns True if the current node or any of its descendants contain a 1.
2. Base Case:
   • If the current node is None, return False because a null node cannot contain a 1.
3. Recursive Case:
   • Recursively check both left and right subtrees.
   • Based on the results, prune subtrees that do not contain a 1 and return whether the current subtree contains a 1.
4. Return the pruned tree:
   • Call the contains_one function from the root.
   • If the root itself is pruned (doesn’t contain a 1), return None.

Time Complexity

• Time Complexity: O(n)
  • We traverse each node exactly once to check and prune, where n is the number of nodes in the binary tree.
• Space Complexity: O(h)
  • The recursion depth will be as deep as the height h of the binary tree. In the worst case (e.g., a skewed tree), the height could be n. In a balanced tree, the height would be log(n).

Hence, the solution efficiently prunes the tree with respect to both time and space.

Try our interview co-pilot at AlgoAdvance.com

• Got blindsided by a question you didn’t expect?

• Spend too much time studying?

• Or simply don’t have the time to go over all 3000 questions?