Задача: 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) для корректировки суммы расстояний для каждого узла на основе суммы расстояний его родительского узла.

😎 Решение:

#include <vector>
#include <unordered_set>
using namespace std;

class Solution {
public:
    vector<int> ans, count;
    vector<unordered_set<int>> graph;
    int N;

    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        this->N = N;
        graph.resize(N);
        ans.resize(N);
        count.resize(N, 1);

        for (auto& edge : edges) {
            graph[edge[0]].insert(edge[1]);
            graph[edge[1]].insert(edge[0]);
        }

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

    void dfs(int node, int parent) {
        for (int child : graph[node]) {
            if (child != parent) {
                dfs(child, node);
                count[node] += count[child];
                ans[node] += ans[child] + count[child];
            }
        }
    }

    void dfs2(int node, int parent) {
        for (int child : graph[node]) {
            if (child != parent) {
                ans[child] = ans[node] - count[child] + N - count[child];
                dfs2(child, node);
            }
        }
    }
};

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