cpp: add «834. Sum of Distances in Tree»

Signed-off-by: Matej Focko <me@mfocko.xyz>

cpp/sum-of-distances-in-tree.cpp

@@ -0,0 +1,66 @@
+#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;
+    }
+};