diff --git a/cpp/sum-of-distances-in-tree.cpp b/cpp/sum-of-distances-in-tree.cpp new file mode 100644 index 0000000..16c62ef --- /dev/null +++ b/cpp/sum-of-distances-in-tree.cpp @@ -0,0 +1,66 @@ +#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; + } +};