Problem Description

You are given the root of a binary tree. Your task is to remove all subtrees that don't contain the value 1.

A subtree consists of a node and all of its descendants. If a subtree doesn't have any node with value 1, it should be completely removed from the tree.

The key points to understand:

• A subtree includes a node plus every node that descends from it
• Any subtree that doesn't contain at least one 1 should be pruned (removed)
• After pruning, you should return the modified tree

For example, if you have a node with value 0 and all its descendants also have value 0, then this entire subtree should be removed since it contains no 1.

The solution works recursively from bottom to top:

1. First, it processes the left and right subtrees by recursively pruning them
2. After pruning children, it checks if the current node should be removed
3. A node is removed if its value is 0 and it has no remaining children (both left and right are null)
4. The condition root.left == root.right cleverly checks if both children are null (since null == null is true)

This bottom-up approach ensures that we first clean up the descendants before deciding whether to keep the current node, effectively removing all subtrees that don't contain any 1.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: A binary tree is a special type of graph where each node has at most two children and there are no cycles.

Is it a tree?

• Yes: The problem explicitly states we're working with a binary tree structure.

DFS

• We arrive at DFS (Depth-First Search) as the recommended approach.

Conclusion: The flowchart suggests using DFS (Depth-First Search) for this binary tree pruning problem.

This makes perfect sense because:

1. We need to traverse the entire tree to check every subtree
2. We need to process children before deciding whether to keep or remove a parent node (post-order traversal)
3. DFS naturally handles the recursive nature of checking each subtree from bottom to top
4. The problem requires examining all descendants of a node before making a decision about that node, which aligns perfectly with DFS's recursive traversal pattern

The solution implements DFS through recursion, where we:

• Recursively process left and right subtrees first
• Make decisions about the current node based on its value and the results from its children
• Return None to prune a subtree or return the node to keep it

Intuition

The key insight is that we need to make decisions about parent nodes based on information from their children. This naturally suggests a bottom-up approach where we process leaves first and work our way up to the root.

Think about it this way: How can we know if a subtree contains a 1? We need to check:

1. Does the current node have value 1?
2. Does any node in the left subtree have value 1?
3. Does any node in the right subtree have value 1?

If the answer to all three questions is "no", then we should remove this entire subtree.

This leads us to a recursive strategy where we first prune the children, then decide about the current node. By processing children first (post-order traversal), we ensure that when we examine a node, its children have already been cleaned up.

The clever part of the solution is recognizing that after pruning:

• If a child subtree doesn't contain any 1, it becomes None
• A node with value 0 and both children as None should also be removed
• The condition root.left == root.right elegantly checks if both are None

This way, the pruning propagates upward: if we remove all descendants of a node with value 0, that node itself becomes a leaf with value 0, which should also be removed. The recursion naturally handles this cascading effect, pruning entire branches that don't contain any 1 values.

Learn more about Tree, Depth-First Search and Binary Tree patterns.

Solution Approach

The solution implements a recursive DFS approach with post-order traversal. Let's walk through the implementation step by step:

Base Case:

if root is None:
    return root

If we encounter a None node (empty tree or reached beyond a leaf), we simply return None. This handles the edge case and serves as our recursion termination condition.

Recursive Pruning:

root.left = self.pruneTree(root.left)
root.right = self.pruneTree(root.right)

We recursively prune both subtrees first. This is crucial because we need to know the final state of the children before deciding about the current node. The pruned subtrees are reassigned to the current node's left and right pointers.

Decision Logic:

if root.val == 0 and root.left == root.right:
    return None
return root

After processing children, we check if the current node should be removed:

• root.val == 0: The node has value 0
• root.left == root.right: Both children are None (since None == None evaluates to True)

If both conditions are met, the node is a leaf with value 0 (or became a leaf after pruning its children), so we return None to remove it. Otherwise, we return the node itself.

Why This Works: The post-order traversal ensures that we process nodes from bottom to top. When a subtree doesn't contain any 1:

1. All its leaves with value 0 are removed first
2. This makes their parents become new leaves
3. If these parents also have value 0, they get removed too
4. This cascades upward until we either reach a node with value 1 or remove the entire subtree

The algorithm runs in O(n) time where n is the number of nodes, as we visit each node exactly once. The space complexity is O(h) where h is the height of the tree, due to the recursion call stack.

Example Walkthrough

Let's walk through a concrete example to see how the pruning algorithm works:

Initial Tree:

        0
       / \
      0   1
     / \
    0   0

We'll trace through the recursive calls step by step:

Step 1: Start at root (value 0)

• Before deciding about the root, we need to prune its children
• Recursively call pruneTree(root.left) where left child has value 0

Step 2: Process left subtree (value 0)

• Before deciding about this node, prune its children
• Recursively call on its left child (value 0, leaf)

Step 3: Process leftmost leaf (value 0)

• Base case: no children to process
• Check: val == 0 ✓ and left == right (both None) ✓
• Return None (remove this node)

Step 4: Back to left subtree, process its right child (value 0, leaf)

• Base case: no children to process
• Check: val == 0 ✓ and left == right (both None) ✓
• Return None (remove this node)

Step 5: Back to left subtree node

• Both children have been pruned and returned None
• Now root.left = None and root.right = None
• Check: val == 0 ✓ and left == right (both None) ✓
• Return None (remove entire left subtree)

Step 6: Back to root, process right child (value 1)

• This is a leaf with value 1
• Check: val == 0 ✗ (value is 1)
• Return the node itself (keep it)

Step 7: Back to root

• Left child has been pruned to None
• Right child remains (value 1)
• Check: val == 0 ✓ but left == right ✗ (None ≠ node with value 1)
• Return root (keep it because it has a valid right child)

Final Tree:

    0
     \
      1

The algorithm successfully removed the entire left subtree (which contained only 0s) while preserving the path to the node with value 1. The root with value 0 is kept because it's needed to maintain the connection to the right subtree containing 1.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

The algorithm performs a post-order traversal of the binary tree, visiting each node exactly once. At each node, it performs constant time operations:

• Checking if the node is null
• Recursively processing left and right subtrees
• Checking if root.val == 0 and root.left == root.right (checking if both children are None)
• Returning either None or the node itself

Since we visit all n nodes exactly once and perform O(1) work at each node, the total time complexity is O(n).

Space Complexity: O(n)

The space complexity comes from the recursive call stack. In the worst case, when the tree is completely skewed (essentially a linked list), the recursion depth can reach n, requiring O(n) space on the call stack.

For a balanced binary tree, the recursion depth would be O(log n), but since we must consider the worst-case scenario, the space complexity is O(n).

Note: The clever condition root.left == root.right checks if both children are None simultaneously, since if both are None, they would be equal. This is equivalent to checking root.left is None and root.right is None.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Order of Operations - Checking Node Before Pruning Children

A common mistake is to check whether the current node should be removed before recursively pruning its children. This leads to incorrect results because the decision about keeping a node depends on the final state of its children after pruning.

Incorrect Implementation:

def pruneTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
    if root is None:
        return None
  
    # WRONG: Checking before pruning children
    if root.val == 0 and not root.left and not root.right:
        return None
  
    # Children are pruned after the check
    root.left = self.pruneTree(root.left)
    root.right = self.pruneTree(root.right)
  
    return root

Why This Fails: Consider a tree where a node with value 0 has children that will be pruned:

    0
   / \
  0   0

The incorrect approach would keep the root node because it initially has children, even though those children should be removed, making the root a leaf with value 0 that should also be removed.

Solution: Always prune children first (post-order traversal), then make the decision about the current node:

def pruneTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
    if root is None:
        return None
  
    # CORRECT: First prune the children
    root.left = self.pruneTree(root.left)
    root.right = self.pruneTree(root.right)
  
    # Then check if current node should be removed
    if root.val == 0 and root.left is None and root.right is None:
        return None
  
    return root

2. Forgetting to Reassign the Pruned Children

Another pitfall is calling the recursive function on children but not reassigning the results back to the parent's pointers.

Incorrect Implementation:

def pruneTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
    if root is None:
        return None
  
    # WRONG: Not reassigning the pruned results
    self.pruneTree(root.left)
    self.pruneTree(root.right)
  
    if root.val == 0 and root.left is None and root.right is None:
        return None
  
    return root

Why This Fails: The recursive calls might return None (indicating the subtree should be removed), but without reassignment, the parent node still points to the original children. This leaves nodes that should be pruned still connected to the tree.

Solution: Always reassign the results of recursive calls:

root.left = self.pruneTree(root.left)
root.right = self.pruneTree(root.right)

3. Misunderstanding the Pruning Condition

Some might think a node should be removed if it has value 0, regardless of its children:

Incorrect Logic:

if root.val == 0:
    return None

Why This Fails: A node with value 0 should be kept if it has at least one child remaining after pruning (meaning there's a 1 somewhere in its subtree). Only leaf nodes with value 0 should be removed.

Solution: Check both conditions: the node has value 0 AND it has no children (is a leaf):

if root.val == 0 and root.left is None and root.right is None:
    return None