67 lines
1.3 KiB
C++
67 lines
1.3 KiB
C++
|
#include <unordered_map>
|
||
|
|
#include <vector>
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
using graph_t = std::unordered_map<int, std::vector<int>>;
|
||
|
|
|
||
|
|
graph_t make_graph(const std::vector<std::vector<int>> &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<int> count;
|
||
|
|
std::vector<int> 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<int>
|
||
|
|
sumOfDistancesInTree(int n, const std::vector<std::vector<int>> &edges) {
|
||
|
|
state s(n);
|
||
|
|
|
||
|
|
auto g = make_graph(edges);
|
||
|
|
dfs_sub(g, s, 0, -1);
|
||
|
|
dfs(g, s, 0, -1);
|
||
|
|
|
||
|
|
return s.total;
|
||
|
|
}
|
||
|
|
};
|