Given the root of a binary tree, replace the value of each node in the tree with the sum of all its cousins' values.
Two nodes of a binary tree are cousins if they have the same depth with different parents.
Return the root of the modified tree.
Note that the depth of a node is the number of edges in the path from the root node to it.
Example 1:
Input: root = [5,4,9,1,10,null,7]
Output: [0,0,0,7,7,null,11]
Explanation: The diagram above shows the initial binary tree and the binary tree after changing the value of each node.
- Node with value 5 does not have any cousins so its sum is 0.
- Node with value 4 does not have any cousins so its sum is 0.
- Node with value 9 does not have any cousins so its sum is 0.
- Node with value 1 has a cousin with value 7 so its sum is 7.
- Node with value 10 has a cousin with value 7 so its sum is 7.
- Node with value 7 has cousins with values 1 and 10 so its sum is 11.
Example 2:
Input: root = [3,1,2]
Output: [0,0,0]
Explanation: The diagram above shows the initial binary tree and the binary tree after changing the value of each node.
- Node with value 3 does not have any cousins so its sum is 0.
- Node with value 1 does not have any cousins so its sum is 0.
- Node with value 2 does not have any cousins so its sum is 0.
Constraints:
The number of nodes in the tree is in the range [1, 105].
1 <= Node.val <= 104
Solutions
Solution 1: Two DFS Traversals
We create a list $s$ to record the sum of the node values at each level of the binary tree, where $s[depth]$ represents the sum of the node values at the $depth$-th level (the root node is at level $0$).
Next, we perform a DFS traversal to calculate the values in the array $s$. Then, we perform another DFS traversal to update the values of each node's children. The value of a child node is equal to the sum of the node values at its level minus the value of the child node and its sibling nodes.
The time complexity is $O(n)$, and the space complexity is $O(n)$. Here, $n$ is the number of nodes in the binary tree.
/** * Definition for a binary tree node. * type TreeNode struct { * Val int * Left *TreeNode * Right *TreeNode * } */funcreplaceValueInTree(root*TreeNode)*TreeNode{s:=[]int{}vardfs1func(*TreeNode,int)dfs1=func(root*TreeNode,depthint){ifroot==nil{return}iflen(s)<=depth{s=append(s,0)}s[depth]+=root.Valdfs1(root.Left,depth+1)dfs1(root.Right,depth+1)}vardfs2func(*TreeNode,int)dfs2=func(root*TreeNode,depthint){l,r:=0,0ifroot.Left!=nil{l=root.Left.Val}ifroot.Right!=nil{r=root.Right.Val}sub:=l+rdepth++ifroot.Left!=nil{root.Left.Val=s[depth]-subdfs2(root.Left,depth)}ifroot.Right!=nil{root.Right.Val=s[depth]-subdfs2(root.Right,depth)}}dfs1(root,0)root.Val=0dfs2(root,0)returnroot}
First, we update the root node's value to $0$, and use a queue $q$ to store all nodes at each level, initially enqueueing the root node.
Then, we traverse the queue, calculate the sum $s$ of all child nodes' values at each level, then calculate the sum $sub$ of each child node and its sibling nodes' values, and then update each child node's value to $s - sub$.
After the traversal ends, we return the root node.
The time complexity is $O(n)$, and the space complexity is $O(n)$. Here, $n$ is the number of nodes in the binary tree.
/** * Definition for a binary tree node. * type TreeNode struct { * Val int * Left *TreeNode * Right *TreeNode * } */funcreplaceValueInTree(root*TreeNode)*TreeNode{root.Val=0q:=[]*TreeNode{root}forlen(q)>0{t:=[]*TreeNode{}s:=0for_,node:=rangeq{ifnode.Left!=nil{t=append(t,node.Left)s+=node.Left.Val}ifnode.Right!=nil{t=append(t,node.Right)s+=node.Right.Val}}for_,node:=rangeq{sub:=0ifnode.Left!=nil{sub+=node.Left.Val}ifnode.Right!=nil{sub+=node.Right.Val}ifnode.Left!=nil{node.Left.Val=s-sub}ifnode.Right!=nil{node.Right.Val=s-sub}}q=t}returnroot}