Binary Tree Pruning

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

Question

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

Return this 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:

Input: {1,#,0,0,1}
Output: {1,#,0,#,1}
Explanation: 
  Only the red nodes satisfy the property "every subtree not containing a 1".
  The diagram on the right represents the answer.


Input: {1,0,1,0,0,0,1}
Output: {1,#,1,#,1}

Input: {1,1,0,1,1,0,1,0}
Output: {1,1,0,1,1,#,1}

Analyzation

Recursion

This is easy to see that if a leaf node which value is 0, will be pruned. That is the only condition to prune the tree.

The Code

class Solution:
    """
    @param root: the root
    @return: the same tree where every subtree (of the given tree) not containing a 1 has been removed
    """
    def pruneTree(self, root):
        # Write your code here
        if root is None:
            return None
        
        # first do left subtree and then right subtree
        root.left = self.pruneTree(root.left)
        root.right = self.pruneTree(root.right)
        
        # if the root is 0 and it has no children, return None to prune the tree
        if root.val == 0 and root.left is None and root.right is None:
            return None
        else:
            return root

Time & Space Complexity

• Time complexity: O(n). Remember that the best case of a Binary Tree is O(logn) since it is a balanced tree but the worst case is O(n) because it can skew to only one side.
• Space complexity: O(n). Remember that the space complexity is always about how many nodes there are in a Binary Tree. If the tree is well-balanced then it is O(logn) but the worst case will be O(n).