#include #include namespace { using graph_t = std::unordered_map>; graph_t make_graph(const std::vector> &edges) { graph_t g; for (const auto &e : edges) { auto [u, v] = std::make_pair(e[0], e[1]); g[u].push_back(v); g[v].push_back(u); } return g; } struct state { int n; std::vector count; std::vector total; state(int n) : n(n), count(n), total(n) {} }; void dfs_sub(graph_t &g, state &s, int u, int parent) { s.count[u] = 1; for (auto v : g[u]) { if (v == parent) { continue; } dfs_sub(g, s, v, u); s.count[u] += s.count[v]; s.total[u] += s.total[v] + s.count[v]; } } void dfs(graph_t &g, state &s, int u, int parent) { for (auto v : g[u]) { if (v == parent) { continue; } s.total[v] = s.total[u] - 2 * s.count[v] + s.n; dfs(g, s, v, u); } } } // namespace struct Solution { std::vector sumOfDistancesInTree(int n, const std::vector> &edges) { state s(n); auto g = make_graph(edges); dfs_sub(g, s, 0, -1); dfs(g, s, 0, -1); return s.total; } };