diff options
Diffstat (limited to 'cousins_binary_tree2.py')
-rw-r--r-- | cousins_binary_tree2.py | 58 |
1 files changed, 58 insertions, 0 deletions
diff --git a/cousins_binary_tree2.py b/cousins_binary_tree2.py new file mode 100644 index 0000000..23e56d5 --- /dev/null +++ b/cousins_binary_tree2.py @@ -0,0 +1,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])) |