Problem Statement in English

You’re given the root of a binary tree. Replace the value of each node in the tree with the sum of the values of its cousins.

Return the root of the modified tree.

Approach

It’s going to be a lot harder to calculate the sum of cousins for each node. Instead we’ll calculate the sum of each level and then update the value of each node with the sum of its level minus the value of its siblings. That’s much easier.

So we have 2 passes, one to calculate the sum of each level and one to update the values of each node.

And we’re done!

Solution in Python


class Solution:
    def replaceValueInTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
        hm = defaultdict(int)

        def scan(node, depth):
            hm[depth] += node.val

            if node.left:
                scan(node.left, depth + 1)
            
            if node.right:
                scan(node.right, depth + 1)

        scan(root, 0)

        def eval(node, offset, depth):
            node.val = hm[depth] - offset

            t = 0
            
            if node.left:
                t += node.left.val
            if node.right:
                t += node.right.val
            
            if node.left:
                eval(node.left, t, depth + 1)
            if node.right:
                eval(node.right, t, depth + 1)

        eval(root, root.val, 0)
    
        return root

Complexity

  • Time: $O(n)$
    Since we scan the tree twice, once to calculate the sum of each level and once to update the values.

  • Space: $O(n)$
    Since we store the sum of each level in a hashmap.

And we are done.