2641. Cousins in Binary Tree II

1class Solution {
2    // Declare a list to store the sum of values at each level.
3    private List<Integer> levelSums = new ArrayList<>();
4
5    // Method to replace values in the tree with the calculated sums.
6    public TreeNode replaceValueInTree(TreeNode root) {
7        calculateLevelSums(root, 0); // Step 1: Calculate sums for all levels.
8        root.val = 0;                // Step 2: Set the root value to 0 as per the requirement.
9        replaceValues(root, 1);      // Step 3: Start replacing the values from the first level.
10        return root;                 // Return the modified tree.
11    }
12
13    // DFS helper method to calculate sum of all nodes at each depth level.
14    private void calculateLevelSums(TreeNode node, int depth) {
15        if (node == null) {
16            return;
17        }
18        // If this is the first time we're visiting this depth level, add a new sum entry.
19        if (levelSums.size() <= depth) {
20            levelSums.add(0);
21        }
22        // Update the current level sum with the value of the current node.
23        levelSums.set(depth, levelSums.get(depth) + node.val);
24        // Recurse on left and right children.
25        calculateLevelSums(node.left, depth + 1);
26        calculateLevelSums(node.right, depth + 1);
27    }
28
29    // DFS helper method to replace node values based on the previously calculated sums.
30    private void replaceValues(TreeNode node, int depth) {
31        if (node == null) {
32            return;
33        }
34        // Compute the values to deduct (current node's children values if they aren't null).
35        int leftValue = node.left == null ? 0 : node.left.val;
36        int rightValue = node.right == null ? 0 : node.right.val;
37
38        // If the left child exists, replace its value based on the sum at the current depth.
39        if (node.left != null) {
40            node.left.val = levelSums.get(depth) - leftValue - rightValue;
41        }
42        // If the right child exists, do the same for the right child.
43        if (node.right != null) {
44            node.right.val = levelSums.get(depth) - leftValue - rightValue;
45        }
46
47        // Recurse on left and right children to continue the process.
48        replaceValues(node.left, depth + 1);
49        replaceValues(node.right, depth + 1);
50    }
51}
52

1#include <vector>
2#include <functional>
3
4// Definition for a binary tree node.
5struct TreeNode {
6    int val;
7    TreeNode *left;
8    TreeNode *right;
9    TreeNode() : val(0), left(nullptr), right(nullptr) {}
10    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
11    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
12};
13
14class Solution {
15public:
16    TreeNode* replaceValueInTree(TreeNode* root) {
17        // Create a vector to store the sum at each depth of the tree.
18        std::vector<int> sumsByDepth;
19
20        // First depth-first search to calculate the sum of values at each depth level.
21        std::function<void(TreeNode*, int)> dfsSum = [&](TreeNode* node, int depth) {
22            if (!node) {
23                return;
24            }
25            // If the current depth is greater than the size of the vector, expand it.
26            if (sumsByDepth.size() <= depth) {
27                sumsByDepth.push_back(0);
28            }
29            // Add the current node's value to the corresponding depth's sum.
30            sumsByDepth[depth] += node->val;
31            // Explore left and right subtrees.
32            dfsSum(node->left, depth + 1);
33            dfsSum(node->right, depth + 1);
34        };
35
36        // Second depth-first search to update the values in the tree.
37        std::function<void(TreeNode*, int)> dfsUpdate = [&](TreeNode* node, int depth) {
38            if (!node) {
39                return;
40            }
41
42            // Store the values for the left and right children if they exist.
43            int leftValue = node->left ? node->left->val : 0;
44            int rightValue = node->right ? node->right->val : 0;
45          
46            // Update the left child's value if it exists.
47            if (node->left) {
48                node->left->val = sumsByDepth[depth] - leftValue - rightValue;
49            }
50            // Update the right child's value if it exists.
51            if (node->right) {
52                node->right->val = sumsByDepth[depth] - leftValue - rightValue;
53            }
54            // Explore left and right subtrees.
55            dfsUpdate(node->left, depth + 1);
56            dfsUpdate(node->right, depth + 1);
57        };
58
59        // Start the first DFS from the root at depth 0.
60        dfsSum(root, 0);
61        // Update the root's value to 0 according to the problem requirement.
62        root->val = 0;
63        // Start the second DFS from the root's children at depth 1.
64        dfsUpdate(root, 1);
65
66        // Return the modified tree root.
67        return root;
68    }
69};
70

1// A tree node may contain a value and pointers to left and right child nodes.
2interface TreeNode {
3    val: number;
4    left: TreeNode | null;
5    right: TreeNode | null;
6}
7
8// This function performs a depth-first traversal to collect the sum of values at each depth.
9function collectSumByDepth(root: TreeNode | null, depth: number, sumAtDepth: number[]): void {
10    if (!root) return;
11  
12    // Ensure the array is large enough to hold the sum for the current depth.
13    if (sumAtDepth.length <= depth) {
14        sumAtDepth.push(0);
15    }
16
17    // Add the current node's value to the sum corresponding to its depth.
18    sumAtDepth[depth] += root.val;
19
20    // Traverse the left and right children.
21    collectSumByDepth(root.left, depth + 1, sumAtDepth);
22    collectSumByDepth(root.right, depth + 1, sumAtDepth);
23}
24
25// This function replaces node values with the collected sum at their respective depth.
26function replaceNodeValueByDepth(root: TreeNode | null, depth: number, sumAtDepth: number[]): void {
27    if (!root) return;
28
29    // The sum of the values of the children nodes.
30    const childSum = (root.left?.val ?? 0) + (root.right?.val ?? 0);
31
32    // Replace the value of the left child if it exists.
33    if (root.left) {
34        root.left.val = sumAtDepth[depth] - childSum;
35    }
36
37    // Replace the value of the right child if it exists.
38    if (root.right) {
39        root.right.val = sumAtDepth[depth] - childSum;
40    }
41  
42    // Continue the traversal for left and right children.
43    replaceNodeValueByDepth(root.left, depth + 1, sumAtDepth);
44    replaceNodeValueByDepth(root.right, depth + 1, sumAtDepth);
45}
46
47// The main function that initiates the process of replacing node values by depth sum.
48function replaceValueInTree(root: TreeNode | null): TreeNode | null {
49    // An array to store the sum of values at each depth.
50    const sumAtDepth: number[] = [];
51
52    // First depth-first search to collect sums at each depth.
53    collectSumByDepth(root, 0, sumAtDepth);
54
55    // Set the root's value to 0 as per instructions.
56    if (root) {
57        root.val = 0;
58    }
59
60    // Second depth-first search to replace the node values by the collected sums at each depth.
61    replaceNodeValueByDepth(root, 1, sumAtDepth);
62
63    return root;
64}
65