Binary Tree Pruning

Problem Statement

Given a binary tree wherein each node either contains the value 0 or 1, the task is to modify the tree so that any subtree that does not have at least one node with the value 1 is completely removed from the original tree. Here, the "subtree" of any node includes the node itself and all of its descendants. The goal is to prune the tree accordingly and return the modified tree which retains subtrees containing the value 1.

Examples

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.

Approach and Intuition

The solution to the problem revolves primarily around a recursive traversal of the binary tree. Let's break down the intuition and approach step by step:

1. Utilize a recursive function that will traverse down from the root to the leaves of the binary tree.
2. For each node, recursively process its left and right children first. This is done so you can decide on the presence or absence of the 1 value in the subtrees originating from these children first before processing the current node.
3. After processing the left and right children:
   • If the left child subtree does not contain a 1, detach (or prune) the left child from the current node.
   • Similarly, if the right child subtree does not contain a 1, detach the right child from the current node.
4. To make prunings easier and decisions straightforward:
   • Let the recursive function return a boolean indicating whether there is at least one 1 in the subtree from the current node downwards.
5. Finally, the result of the recursive function for the root will give the necessary modifications directly, as the whole tree would have been pruned from the bottom up, retaining only the necessary parts that contain at least one node with the value 1.

This recursive algorithm efficiently prunes the tree in a depth-first manner, ensuring that every node is visited and appropriately retained or discarded based on the presence of the 1 in its subtree.

Solutions

class Solution {
    public TreeNode removeZeros(TreeNode root) {
        return hasOne(root) ? root : null;
    }
    
    public boolean hasOne(TreeNode curr) {
        if (curr == null) return false;
    
        boolean leftHasOne = hasOne(curr.left);
        boolean rightHasOne = hasOne(curr.right);
    
        if (!leftHasOne) curr.left = null;
        if (!rightHasOne) curr.right = null;
    
        return curr.val == 1 || leftHasOne || rightHasOne;
    }
}