Problem statement

https://leetcode.com/problems/binary-tree-pruning/

Solution

Use dfs(node) function which returns True only if all subtree has only zeroes. Evaluate it for left and right children: if for any of them answer is True, we cut this children. Also we return True if answer for both child is true and value of node is equal to zero. In the end, we check if answer is True and if it is, we return empty tree end if it is not, we return original root.

Complexity

Time complexity is O(n), space complexity is O(h).

Code

class Solution:
    def pruneTree(self, root):
        def dfs(node):
            if not node: return True
            left, right = dfs(node.left), dfs(node.right)
            if left: node.left = None
            if right: node.right = None
            return left and right and node.val == 0
        
        return root if not dfs(root) else None