Задача: 834. Sum of Distances in Tree
Сложность: hard

Есть неориентированное связное дерево с n узлами, пронумерованными от 0 до n - 1, и n - 1 ребрами.

Вам даны целое число n и массив edges, где edges[i] = [ai, bi] указывает, что существует ребро между узлами ai и bi в дереве.

Верните массив answer длиной n, где answer[i] — это сумма расстояний между i-м узлом в дереве и всеми другими узлами.

Пример:

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.

👨‍💻 Алгоритм:

1⃣Постройте дерево и выполните обход в глубину (DFS) для расчета количества узлов в поддереве и суммы расстояний до всех узлов поддерева.

2⃣На основе полученных данных рассчитайте сумму расстояний для корня дерева.

3⃣Выполните второй обход в глубину (DFS) для корректировки суммы расстояний для каждого узла на основе суммы расстояний его родительского узла.

😎 Решение:

class Solution {
    private $ans;
    private $count;
    private $graph;
    private $N;

    function sumOfDistancesInTree($N, $edges) {
        $this->N = $N;
        $this->graph = array_fill(0, $N, []);
        $this->ans = array_fill(0, $N, 0);
        $this->count = array_fill(0, $N, 1);

        foreach ($edges as $edge) {
            $this->graph[$edge[0]][] = $edge[1];
            $this->graph[$edge[1]][] = $edge[0];
        }

        $this->dfs(0, -1);
        $this->dfs2(0, -1);
        return $this->ans;
    }

    function dfs($node, $parent) {
        foreach ($this->graph[$node] as $child) {
            if ($child != $parent) {
                $this->dfs($child, $node);
                $this->count[$node] += $this->count[$child];
                $this->ans[$node] += $this->ans[$child] + $this->count[$child];
            }
        }
    }

    function dfs2($node, $parent) {
        foreach ($this->graph[$node] as $child) {
            if ($child != $parent) {
                $this->ans[$child] = $this->ans[$node] - $this->count[$child] + $this->N - $this->count[$child];
                $this->dfs2($child, $node);
            }
        }
    }
}

Ставь 👍 и забирай 📚 Базу знаний