2641. Cousins in Binary Tree II


Problem Description

In this problem, we are given the root of a binary tree. A binary tree is a data structure where each node has at most two children, referred to as the left child and the right child. We are asked to modify the tree by replacing each node's value with the sum of the values of all its cousin nodes. Cousins are defined as nodes that are at the same level (have the same depth) but do not share the same parent. The depth of a node is measured by the number of edges from the root node to the given node.

The challenge is to perform this transformation in place, that is, without creating a new tree but rather by modifying the values of nodes in the original tree.

Intuition

The solution to this problem involves two main steps:

  1. Calculate the sum of values at each depth in the tree.
  2. Replace the value of each node with the sum of its cousin's values, which is the total sum at its depth minus the values of the node's children (since children are not considered cousins).

We can achieve this through two depth-first search (DFS) traversals:

  • The first DFS (dfs1) traverses the tree to calculate the sum of values at each depth. We maintain an array s where s[d] represents the sum of values of all nodes at depth d. As we recurse, we pass the depth d and increment it when moving to the children.

  • The second DFS (dfs2) uses the computed sums to replace each node's value with the correct sum of its cousins' values. For each node, we subtract the values of its children from the sum at its depth s[d] (if it has children) and assign this value to the node's children. We do not need to modify the value of the root node as it has no cousins.

The result is an updated tree where each node's value represents the sum of all its cousins' values, as required by the problem statement.

Learn more about Tree, Depth-First Search, Breadth-First Search and Binary Tree patterns.

Solution Approach

The implementation of this problem can be broken down as follows:

  • Firstly, we define two recursive functions dfs1 and dfs2 aimed at conducting depth-first search traversals. Depth-First Search (DFS) is an algorithm for traversing or searching tree or graph data structures. It starts at the root node and explores as far as possible along each branch before backtracking.

  • The dfs1 function is responsible for calculating the sum of values at each depth of the tree. It accepts two parameters: root, which is the current node being visited, and d, which represents the current depth. We utilize a list s where each s[d] corresponds to the sum of node values at depth d. When the function visits a node, it adds the node's value to s[d] and recursively calls itself for the left and right children, with the depth d increased by 1.

  • After the first traversal with dfs1, we have the sum of values for each level of the tree stored in s. However, to find the sum of cousins, we must subtract the values of any children a node might have, as they are not its cousins.

  • To carry out this subtraction and update the node values accordingly, we use the dfs2 function. Like dfs1, dfs2 also takes root and d as parameters. For each node, we calculate the sum of its children's values (if any) and subtract this sum from s[d] to find the sum of its cousins' values. We then recursively call dfs2 for the left and right children, again increasing the depth d by 1.

  • Note that we don't change the value of the root node itself, as it doesn't have any cousins. Therefore, we set root.val to 0 before calling dfs2.

  • After the second traversal with dfs2, all nodes will contain the sum of their cousins' values.

  • Lastly, we return the modified root of the tree.

The algorithm uses the DFS pattern to explore the nodes. It also makes use of recursion to navigate through the tree from the root to the leaves efficiently. By dividing the problem into two parts, we can achieve the solution without any need for complex data structures or additional storage beyond the call stack and the sum array s.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Breadth first search can be used to find the shortest path between two nodes in a directed graph.

Example Walkthrough

Let’s consider the following binary tree for our example:

1     1
2   /   \
3  2     3
4 / \   /
54   5 6

For simplicity, we’ll refer to each node by its initial value.

Step 1: Run dfs1 starting with the root node (node 1) at depth 0. This will calculate the sum of all node values at each depth.

  • At depth 0, s[0] becomes 1 (value of root node).
  • At depth 1, s[1] becomes 2 + 3 = 5 (values of node 2 and node 3).
  • At depth 2, s[2] becomes 4 + 5 + 6 = 15 (values of node 4, node 5, and node 6).

Step 2: Before running dfs2, set root.val to 0 since the root has no cousins.

Step 3: Run dfs2 starting with node 1 (which now has a value of 0) at depth 0. This traversal will update each node with the sum of the cousins' values.

  • Node 1's value remains 0.
  • Node 2 is at depth 1, and its children's values sum up to 4 + 5 = 9. We calculate its new value as s[1] - 9 = 5 - 9 = -4.
  • Node 3 is also at depth 1, but its child's value is 6. Its new value will be s[1] - 6 = 5 - 6 = -1.
  • Node 4 is at depth 2, with no children. Its new value will be s[2] = 15.
  • Node 5 is at depth 2, with no children. Its new value will also be s[2] = 15.
  • Node 6 is at depth 2, with no children. Its new value will be s[2] = 15, just like nodes 4 and 5.

After running both traversals, the modified tree will look like:

1     0
2   /   \
3 -4    -1
4 / \    /
515  15 15

Note that the sum values for the cousins are negative for nodes 2 and 3 because their children's values in the original tree exceeded the sum of values at that depth. This occurrence is based on the instructions given in the content and demonstrates that the algorithm processes the sum of node values at each depth and then properly subtracts the children's values to compute the final update during the second DFS traversal.

Solution Implementation

1# Definition for a binary tree node.
2class TreeNode:
3    def __init__(self, val=0, left=None, right=None):
4        self.val = val
5        self.left = left
6        self.right = right
7
8class Solution:
9    def replaceValueInTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
10        # Helper function to perform DFS and sum up the values at each depth.
11        def dfs_sum_values(node, depth):
12            if node is None:
13                return
14            # If we encounter a new depth, initialize sum for that depth.
15            if len(depth_sums) <= depth:
16                depth_sums.append(0)
17            # Add the current node's value to the corresponding depth's sum.
18            depth_sums[depth] += node.val
19            # Recursive calls for the left and right subtrees, with increased depth.
20            dfs_sum_values(node.left, depth + 1)
21            dfs_sum_values(node.right, depth + 1)
22
23        # Helper function to perform DFS and replace the values of non-root nodes.
24        def dfs_replace_values(node, depth):
25            if node is None:
26                return
27            # Calculate the combined value of the current node's children.
28            children_sum = (node.left.val if node.left else 0) + (node.right.val if node.right else 0)
29            # Replace the values of the left and right children.
30            if node.left:
31                node.left.val = depth_sums[depth] - children_sum
32            if node.right:
33                node.right.val = depth_sums[depth] - children_sum
34            # Recursive calls for the left and right subtrees, with increased depth.
35            dfs_replace_values(node.left, depth + 1)
36            dfs_replace_values(node.right, depth + 1)
37
38        # Initialize the list that will hold the sums of values at each depth.
39        depth_sums = []
40        # First DFS to calculate sums at each depth.
41        dfs_sum_values(root, 0)
42        # Set the root value to 0 since the question seems to be replacing non-root values.
43        root.val = 0
44        # Second DFS to replace values in the tree.
45        dfs_replace_values(root, 1)
46      
47        return root
48
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

Time and Space Complexity

The time complexity for the dfs1 function is O(n) where n is the number of nodes in the binary tree since it visits each node exactly once to calculate the sum at each depth. Similarly, dfs2 also has a time complexity of O(n) as it again visits each node exactly once to replace the values. Hence the combined time complexity of the entire replaceValueInTree method is also O(n).

For space complexity, the code uses O(h) extra space for recursion where h is the height of the binary tree, which in the worst case of a skewed tree can be O(n). Additionally, the list s can grow up to O(h) for storing the sum at each depth, again leading to O(n) in the worst case. Therefore, the overall space complexity is O(n) as well.

Learn more about how to find time and space complexity quickly using problem constraints.


Fast Track Your Learning with Our Quick Skills Quiz:

Which of the following is a good use case for backtracking?


Recommended Readings


Got a question? Ask the Monster Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


🪄