Cousins in binary Tree II

XXX. Replace Values in Tree with Cousin Sums

Tree

Breadth-First Search

Problem Statement:

Given the root of a binary tree, replace the value of each node with the sum of its cousins' values. The root node should have a value of 0, as it has no cousins.

Algorithm:

1. Perform a BFS to calculate the sum of nodes at each level and store these sums in a list.
2. Perform a second BFS to update the value of each node as follows:
   • For each node, calculate the sum of its sibling nodes.
   • Update the node's value to the difference between the total level sum and its sibling sum.

Complexity:

Time Complexity: O(n), where n is the number of nodes in the tree.
Space Complexity: O(w), where w is the maximum width of the tree (queue size in BFS).

Java Implementation:

public TreeNode replaceValueInTree(TreeNode root) {
    if (root == null) return root;

    Queue nodeQueue = new LinkedList<>();
    nodeQueue.offer(root);
    List levelSums = new ArrayList<>();

    // First BFS: Calculate sum of nodes at each level
    while (!nodeQueue.isEmpty()) {
        int levelSum = 0;
        int levelSize = nodeQueue.size();
        for (int i = 0; i < levelSize; ++i) {
            TreeNode currentNode = nodeQueue.poll();
            levelSum += currentNode.val;
            if (currentNode.left != null) nodeQueue.offer(currentNode.left);
            if (currentNode.right != null) nodeQueue.offer(currentNode.right);
        }
        levelSums.add(levelSum);
    }

    // Second BFS: Update each node's value to sum of its cousins
    nodeQueue.offer(root);
    int levelIndex = 1;
    root.val = 0; // Root has no cousins
    while (!nodeQueue.isEmpty()) {
        int levelSize = nodeQueue.size();
        for (int i = 0; i < levelSize; ++i) {
            TreeNode currentNode = nodeQueue.poll();

            int siblingSum =
                (currentNode.left != null ? currentNode.left.val : 0) +
                (currentNode.right != null ? currentNode.right.val : 0);

            if (currentNode.left != null) {
                currentNode.left.val =
                    levelSums.get(levelIndex) - siblingSum;
                nodeQueue.offer(currentNode.left);
            }

            if (currentNode.right != null) {
                currentNode.right.val =
                    levelSums.get(levelIndex) - siblingSum;
                nodeQueue.offer(currentNode.right);
            }
        }
        levelIndex++;
    }

    return root;
}
                

Python Implementation:

from collections import deque

def replace_value_in_tree(root):
    if not root:
        return root

    node_queue = deque([root])
    level_sums = []

    # First BFS: Calculate sum of nodes at each level
    while node_queue:
        level_sum = 0
        level_size = len(node_queue)
        for _ in range(level_size):
            current_node = node_queue.popleft()
            level_sum += current_node.val
            if current_node.left:
                node_queue.append(current_node.left)
            if current_node.right:
                node_queue.append(current_node.right)
        level_sums.append(level_sum)

    # Second BFS: Update each node's value to sum of its cousins
    node_queue.append(root)
    level_index = 1
    root.val = 0  # Root has no cousins
    while node_queue:
        level_size = len(node_queue)
        for _ in range(level_size):
            current_node = node_queue.popleft()

            sibling_sum = (
                (current_node.left.val if current_node.left else 0) +
                (current_node.right.val if current_node.right else 0)
            )

            if current_node.left:
                current_node.left.val = level_sums[level_index] - sibling_sum
                node_queue.append(current_node.left)

            if current_node.right:
                current_node.right.val = level_sums[level_index] - sibling_sum
                node_queue.append(current_node.right)

        level_index += 1

    return root
                

C++ Implementation:

#include 
#include 
using namespace std;

TreeNode* replaceValueInTree(TreeNode* root) {
    if (!root) return root;

    queue nodeQueue;
    nodeQueue.push(root);
    vector levelSums;

    // First BFS: Calculate sum of nodes at each level
    while (!nodeQueue.empty()) {
        int levelSum = 0;
        int levelSize = nodeQueue.size();
        for (int i = 0; i < levelSize; ++i) {
            TreeNode* currentNode = nodeQueue.front();
            nodeQueue.pop();
            levelSum += currentNode->val;
            if (currentNode->left) nodeQueue.push(currentNode->left);
            if (currentNode->right) nodeQueue.push(currentNode->right);
        }
        levelSums.push_back(levelSum);
    }

    // Second BFS: Update each node's value to sum of its cousins
    nodeQueue.push(root);
    int levelIndex = 1;
    root->val = 0; // Root has no cousins
    while (!nodeQueue.empty()) {
        int levelSize = nodeQueue.size();
        for (int i = 0; i < levelSize; ++i) {
            TreeNode* currentNode = nodeQueue.front();
            nodeQueue.pop();

            int siblingSum =
                (currentNode->left ? currentNode->left->val : 0) +
                (currentNode->right ? currentNode->right->val : 0);

            if (currentNode->left) {
                currentNode->left->val =
                    levelSums[levelIndex] - siblingSum;
                nodeQueue.push(currentNode->left);
            }

            if (currentNode->right) {
                currentNode->right->val =
                    levelSums[levelIndex] - siblingSum;
                nodeQueue.push(currentNode->right);
            }
        }
        levelIndex++;
    }

    return root;
}