Cousins in Binary Tree II

medium tree bfs level order

Problem

Given the root of a binary tree, replace the value of every node with the sum of the values of all its cousins. Two nodes are cousins if they are on the same level but do not share the same parent. Then return the modified tree.

The cousins sum for a node equals the total of its level minus the values of its own sibling group (the node plus any node sharing its parent). The root has no cousins, so it becomes 0.

Inputroot = [5,4,9,1,10,null,7]

Output[0,0,0,7,7,null,11]

Root → 0. Level total 13 for {4,9} (both children of 5) → each 0. Bottom total 18: 1 and 10 (under 4) → 18 − 11 = 7; node 7 (under 9) → 18 − 7 = 11.

level-order (use "null" for missing children)

from collections import deque

def replace_value_in_tree(root):
    root.val = 0
    queue = deque([root])
    while queue:
        level = list(queue)
        level_sum = 0
        for node in level:
            if node.left:  level_sum += node.left.val
            if node.right: level_sum += node.right.val
        for _ in range(len(queue)):
            node = queue.popleft()
            sibling = (node.left.val if node.left else 0) + \
                      (node.right.val if node.right else 0)
            if node.left:
                node.left.val = level_sum - sibling
                queue.append(node.left)
            if node.right:
                node.right.val = level_sum - sibling
                queue.append(node.right)
    return root

function replaceValueInTree(root) {
  root.val = 0;
  let queue = [root];
  while (queue.length) {
    let levelSum = 0;
    for (const node of queue) {
      if (node.left) levelSum += node.left.val;
      if (node.right) levelSum += node.right.val;
    }
    const next = [];
    for (const node of queue) {
      const sib = (node.left ? node.left.val : 0) +
                  (node.right ? node.right.val : 0);
      if (node.left) { node.left.val = levelSum - sib; next.push(node.left); }
      if (node.right) { node.right.val = levelSum - sib; next.push(node.right); }
    }
    queue = next;
  }
  return root;
}

class Solution {
    public TreeNode replaceValueInTree(TreeNode root) {
        root.val = 0;
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            int size = queue.size();
            int levelSum = 0;
            for (TreeNode node : queue) {
                if (node.left != null) levelSum += node.left.val;
                if (node.right != null) levelSum += node.right.val;
            }
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                int sib = (node.left != null ? node.left.val : 0)
                        + (node.right != null ? node.right.val : 0);
                if (node.left != null) { node.left.val = levelSum - sib; queue.offer(node.left); }
                if (node.right != null) { node.right.val = levelSum - sib; queue.offer(node.right); }
            }
        }
        return root;
    }
}

TreeNode* replaceValueInTree(TreeNode* root) {
    root->val = 0;
    vector<TreeNode*> queue = {root};
    while (!queue.empty()) {
        int levelSum = 0;
        for (auto node : queue) {
            if (node->left) levelSum += node->left->val;
            if (node->right) levelSum += node->right->val;
        }
        vector<TreeNode*> next;
        for (auto node : queue) {
            int sib = (node->left ? node->left->val : 0)
                    + (node->right ? node->right->val : 0);
            if (node->left) { node->left->val = levelSum - sib; next.push_back(node->left); }
            if (node->right) { node->right->val = levelSum - sib; next.push_back(node->right); }
        }
        queue = next;
    }
    return root;
}

Explanation

A node's cousins are everyone on its level except itself and its siblings. So cousinSum = levelTotal − siblingGroupTotal, where the sibling group of a node is all children of the same parent.

The neat trick is to process each level by looking at it from the parents above. When we sit on level L and look at all the children we are about to produce, their combined value is the total of level L+1.

For each parent we know the sum of its own children (the sibling group). The new value every child should take is therefore levelSum(L+1) − thisParentsChildrenSum — the level total with this child's own family removed.

So we BFS level by level: first add up all grandchildren-to-be to get the next level's total, then walk the parents again and assign each child levelSum − siblingSum. The root is set to 0 up front since it has no cousins.

Example on [5,4,9,1,10,null,7]: root → 0. Next level total 4+9 = 13; 4 and 9 are siblings (children of 5), so each becomes 13 − 13 = 0. Bottom level total 1+10+7 = 18; 1 and 10 (children of 4) become 18 − 11 = 7, while 7 (child of 9) becomes 18 − 7 = 11 → [0,0,0,7,7,null,11].

Time: O(n) Space: O(n)