Problem description

We are given the head node root of a binary tree, where additionally every node’s value is either a 0 or a 1.

Return the same tree where every subtree (of the given tree) not containing a 1 has been removed.

(Recall that the subtree of a node X is X, plus every node that is a descendant of X.)

Example 1: Input: [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: [1,0,1,0,0,0,1] Output: [1,null,1,null,1]

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

Note:

The binary tree will have at most 100 nodes. The value of each node will only be 0 or 1.

Solution

1. For each node, we can check whether all nodes of this tree is Zero. If it is, we return make this node None
2. If not, we prune its left and right subtree, and attach the modified subtrees’ root as the child of the current node.

Below is python implementation

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def pruneTree(self, root: TreeNode) -> TreeNode:
        
        
        def allzero(root)->bool:
            if not root:
                return True
            if root.val != 0:
                return False
            return allzero(root.left) and allzero(root.right)
        
        if not root:
            return None
        if allzero(root):
            return None
        left = self.pruneTree(root.left)
        right = self.pruneTree(root.right)
        root.left = left
        root.right = right
        return root