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):
        def containsOne(node):
            if not node:
                return False
            return containsOne(node.left) or containsOne(node.right) or node.val == 1
        if not root:
            return None
        root.left = self.pruneTree(root.left)
        root.right = self.pruneTree(root.right)
        return root if containsOne(root) else None