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

#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;
    }
};