BinarySearch 0080 Tree Pruning

Problem statement

https://binarysearch.com/problems/Tree-Pruning/

Solution

Equal to Leetcode 0814. Binary Tree Pruning

Complexity

It is O(n) for time and O(h) for space.

Code

class Solution:
    def solve(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