# 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]))