Binary Tree Pruning

Updated on 15 May, 2025
Binary Tree Pruning header image

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

  • Java
  • Python
java
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;
    }
}

The provided Java solution addresses the problem of pruning a binary tree by removing all subtrees that do not contain the value 1. The Solution class includes two methods: removeZeros and hasOne.

  • removeZeros(TreeNode root):

    • Takes the root of a binary tree as its parameter.
    • Utilizes the hasOne method to determine if the binary tree rooted at the specified node should be kept or set to null based on whether it contains a 1.
    • Returns the original root if it contains a 1; otherwise, it returns null.
  • hasOne(TreeNode curr):

    • Determines recursively if the tree rooted at curr contains at least one node with the value 1.
    • Checks both the left and right subtrees for the presence of the number 1.
    • If either subtree does not contain a 1, it sets that subtree reference in the current node to null, effectively pruning the tree.
    • Returns true if the current node itself contains a 1 or if any of its subtrees contain a 1, and false otherwise.

This strategy efficiently prunes all the subtrees that do not contribute to containing the value 1, ensuring the resulting binary tree is minimal based on the specified criterion. The recursive approach allows for thorough traversal and pruning of the tree in a depth-first manner.

python
class Solution:
    def pruneBinaryTree(self, root: TreeNode) -> TreeNode:
        
        def hasOne(node: TreeNode) -> bool:
            if not node:
                return False
            
            left_has_one = hasOne(node.left)
            right_has_one = hasOne(node.right)
            
            if not left_has_one:
                node.left = None
                
            if not right_has_one:
                node.right = None
                
            return node.val == 1 or left_has_one or right_has_one

        return root if hasOne(root) else None

This Python solution involves pruning a binary tree such that every subtree of the tree not containing a 1 is removed. This process is conducted in a manner that leaves only the branches of the tree that have at least one node with the value 1.

The solution defines a class Solution with the method pruneBinaryTree, which takes the binary tree's root node as an argument and returns the pruned tree's root node. The method utilizes a nested helper function hasOne to determine if a subtree contains the value 1.

  • The hasOne function:
    • Returns False immediately if the current node is None.
    • Recursively checks both the left and right children of the current node to see if they contain a 1.
    • Removes the left child if it does not contain a 1 by setting node.left to None.
    • Performs a similar operation on the right child.
    • Returns True if the current node's value is 1 or if any of its children contain a 1, indicating the subtree should be kept.

Finally, the pruneBinaryTree method:

  • Calls hasOne on the root node.
  • Returns the root node if hasOne evaluates to True.
  • Returns None if the root subtree does not contain any nodes with value 1, effectively removing all nodes from the tree.

Use this pruning tactic to efficiently clean your binary tree, focusing only on the parts of the tree that meet a specific numeric condition.

Comments

No comments yet.