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
|
# 2641. Cousins in Binary Tree II
from collections import Counter
"""
given the 'root' of a binary tree, replace the value of each node in the tree
with the sum of all its cousins' values. two nodes of a binary tree are
cousins if they have the same depth with different paths. return 'root' of
the modified tree. note that the depth of a nodie is a number of edges in the
path from the root node to it.
"""
# Definition for a binary tree node.
class TreeNode(object):
def __init__(self, val=0, left=None, right=None):
self.val = val
self.left = left
self.right = right
class Solution(object):
def replaceValueInTree(self, root):
"""
:type root: Optional[TreeNode]
:rtype: Optional[TreeNode]
"""
sum = Counter()
def dfs1(r, l):
if not r:
return
sum[l] += r.val
dfs1(r.left, l + 1)
dfs1(r.right, l + 1)
dfs1(root, 0)
def dfs2(r, l, curr):
cousin_sum = (
sum[l + 1]
- (r.left.val if r.left else 0)
- (r.right.val if r.right else 0)
)
if r.left:
curr.left = TreeNode(cousin_sum)
dfs2(r.left, l + 1, curr.left)
if r.right:
curr.right = TreeNode(cousin_sum)
dfs2(r.right, l + 1, curr.right)
return curr
return dfs2(root, 0, TreeNode(0))
if __name__ == "__main__":
obj = Solution()
print(obj.replaceValueInTree(root=[5, 4, 9, 1, 10, None, 7]))
print(obj.replaceValueInTree(root=[3, 1, 2]))
|