834. Sum of Distances in Tree

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 a_i and b_i 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

• LinkedIn
• GitHub