leetcode-solution/Python/2641-cousins-in-binary-tree-ii.py at main · hogan-tech/leetcode-solution · 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
# time complexity: O(n)
# space complexity: O(n)

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def __init__(self):
        self.levelSums = [0] * 100000

    def replaceValueInTree(self, root):
        self.calculateLevelSum(root, 0)
        self.replaceValueInTreeInternal(root, 0, 0)
        return root

    def calculateLevelSum(self, node, level):
        if node is None:
            return
        self.levelSums[level] += node.val
        self.calculateLevelSum(node.left, level + 1)
        self.calculateLevelSum(node.right, level + 1)

    def replaceValueInTreeInternal(self, node, siblingSum, level):
        if node is None:
            return
        leftChildVal = 0 if node.left is None else node.left.val
        rightChildVal = 0 if node.right is None else node.right.val

        if level == 0 or level == 1:
            node.val = 0
        else:
            node.val = self.levelSums[level] - node.val - siblingSum
        self.replaceValueInTreeInternal(
            node.left, rightChildVal, level + 1
        )
        self.replaceValueInTreeInternal(
            node.right, leftChildVal, level + 1
        )


root = TreeNode(5)
root.left = TreeNode(4)
root.right = TreeNode(9)
root.left.left = TreeNode(1)
root.left.right = TreeNode(10)
root.right.right = TreeNode(7)
print(Solution().replaceValueInTree(root))