834. Sum of Distances in Tree
Hard
There is an undirected connected tree with n
nodes labeled from 0
to n - 1
and n - 1
edges.
You are given the integer n
and the array edges
where edges[i] = [ai, bi]
indicates that there is an edge between nodes ai and bi in the tree.
Return an array answer
of length n
where answer[i]
is the sum of the distances between the ith
node in the tree and all other nodes.
Example 1:
- Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
- Output: [8,12,6,10,10,10]
- Explanation: The tree is shown above. We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) equals 1 + 1 + 2 + 2 + 2 = 8. Hence, answer[0] = 8, and so on.
Example 2:
- Input: n = 1, edges = []
- Output: [0]
Example 3:
- Input: n = 2, edges = [[1,0]]
- Output: [1,1]
Constraints:
1 <= n <= 3 * 104
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
ai != bi
- The given input represents a valid tree.
Solution:
class Solution {
/**
* @param Integer $n
* @param Integer[][] $edges
* @return Integer[]
*/
function sumOfDistancesInTree($n, $edges) {
$g = array_fill(0, $n, array());
foreach ($edges as $e) {
$a = $e[0];
$b = $e[1];
array_push($g[$a], $b);
array_push($g[$b], $a);
}
$ans = array_fill(0, $n, 0);
$size = array_fill(0, $n, 0);
$dfs1 = function ($i, $fa, $d) use (&$ans, &$size, &$g, &$dfs1) {
$ans[0] += $d;
$size[$i] = 1;
foreach ($g[$i] as &$j) {
if ($j != $fa) {
$dfs1($j, $i, $d + 1);
$size[$i] += $size[$j];
}
}
};
$dfs2 = function ($i, $fa, $t) use (&$ans, &$size, &$g, &$dfs2, $n) {
$ans[$i] = $t;
foreach ($g[$i] as &$j) {
if ($j != $fa) {
$dfs2($j, $i, $t - $size[$j] + $n - $size[$j]);
}
}
};
$dfs1(0, -1, 0);
$dfs2(0, -1, $ans[0]);
return $ans;
}
}
Contact Links
If you found this series helpful, please consider giving the repository a star on GitHub or sharing the post on your favorite social networks 😍. Your support would mean a lot to me!