Insufficient Nodes in Root to Leaf Paths LeetCode solution in Python

1
class Solution:
2
    """
3
    we can try to solve this problem using depth first traversal
4
    """
5

6
    def dfs(self, root, sum_so_far, limit):
7
        if root is None:
8
            return None, 0
9

10
        x, left = self.dfs(root.left, sum_so_far + root.val, limit)
11
        y, right = self.dfs(root.right, sum_so_far + root.val, limit)
12
        # print('for node= {}, sum_so_far= {}, left= {}, right= {}'.format(root.val, sum_so_far, left, right))
13
        if root.left is None and root.right is None:
14
            # it is leaf, left and right should be 0
15
            if sum_so_far + root.val < limit:
16
                # node is insufficient
17
                return None, root.val
18
            else:
19
                # node is sufficient
20
                return root, root.val
21
        elif root.left is not None and root.right is None:
22
            root.left = x
23
            if sum_so_far + root.val + left < limit:
24
                # node is insufficient
25
                return None, root.val + left
26
            else:
27
                return root, root.val + left
28

29
        elif root.left is None and root.right is not None:
30
            root.right = y
31
            if sum_so_far + root.val + right < limit:
32
                return None, root.val + right
33
            else:
34
                return root, root.val + right
35

36
        elif root.left is not None and root.right is not None:
37
            root.left = x
38
            root.right = y
39
            if (
40
                sum_so_far + root.val + left < limit
41
                and sum_so_far + root.val + right < limit
42
            ):
43
                return None, max(root.val + left, root.val + right)
44
            elif (
45
                sum_so_far + root.val + left < limit
46
                and sum_so_far + root.val + right > limit
47
            ):
48
                return root, root.val + right
49
            elif (
50
                sum_so_far + root.val + left > limit
51
                and sum_so_far + root.val + right < limit
52
            ):
53
                return root, root.val + left
54
            else:
55
                return root, max(root.val + left, root.val + right)
56

57
    def sufficientSubset(
58
        self, root: Optional[TreeNode], limit: int
59
    ) -> Optional[TreeNode]:
60
        root, _ = self.dfs(root, 0, limit)
61
        return root