learning/l33tcode/sum-root-to-leaf-numbers.py at master · fspv/learning · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def sumNumbers(self, root: TreeNode) -> int:
        # O(1) space solution

        def find_predecessor(node):
            path_length = 1
            predecessor = node.left

            while predecessor and predecessor.right and predecessor.right != node:
                path_length += 1
                predecessor = predecessor.right

            return path_length, predecessor

        def preorder_traversal(node):
            total = 0
            cur_sum = 0
            while node:
                path_length, predecessor = find_predecessor(node)
                if predecessor and predecessor.right:
                    # reset sum value
                    if not predecessor.left:  # means leaf node
                        total += cur_sum

                    cur_sum = cur_sum // (10 ** path_length)

                    # reset predecessor ptr and right count
                    predecessor.right = None

                    # go right or back
                    node = node.right
                elif predecessor and not predecessor.right:
                    # visit
                    cur_sum = cur_sum * 10 + node.val

                    # set link and go left
                    predecessor.right = node
                    node = node.left
                elif not predecessor:
                    # visit
                    cur_sum = cur_sum * 10 + node.val

                    if not node.right:
                        # we've reached the rightmost point
                        total += cur_sum

                    # go right or back
                    node = node.right

            return total

        return preorder_traversal(root)

    def sumNumbersONSpace(self, root: TreeNode) -> int:
        def dfs(node: TreeNode, prev: int) -> int:
            if not node:
                return 0

            current = prev * 10 + node.val

            if not (node.left or node.right):
                return current

            total = 0
            if node.left:
                total += dfs(node.left, current)
            if node.right:
                total += dfs(node.right, current)

            return total

        return dfs(root, 0)