Submission #51888815

---

Source Code Expand

#include <bits/stdc++.h>

using namespace std;

using i64 = int64_t;

int main() {
    // in
    int n;
    cin >> n;
    vector<int> a(n - 1), b(n - 1);
    vector<vector<int>> tree(n);
    for (int i = 0; i < n - 1; ++i) {
        cin >> a[i] >> b[i];
        --a[i], --b[i];
        tree[a[i]].push_back(b[i]);
        tree[b[i]].push_back(a[i]);
    }

    vector<i64> c(n);
    for (auto &e : c) {
        cin >> e;
    }

    // solve
    const i64 sum_c = accumulate(c.begin(), c.end(), 0ll);
    const int root = 0;

    // subtree_sum[i] : sum c[x] (x in subtree(i))
    vector<i64> subtree_sum(n);
    auto dfs = [&](auto &&self, const int v, const int par) -> i64 {
        subtree_sum[v] = c[v];
        for (const int to : tree[v]) {
            if (to == par) {
                continue;
            }
            subtree_sum[v] += self(self, to, v);
        }
        return subtree_sum[v];
    };
    dfs(dfs, root, -1);

    auto find_centroid = [&](auto &&self, const int v, const int par) -> int {
        bool is_centroid = true;
        for (const int to : tree[v]) {
            if (to == par) {
                continue;
            }
            const int res = self(self, to, v);
            if (res != -1) {
                return res;
            }

            if (subtree_sum[to] > sum_c / 2) {
                is_centroid = false;
            }
        }
        if ((sum_c - subtree_sum[v]) > sum_c / 2) {
            is_centroid = false;
        }
        return is_centroid ? v : -1;
    };
    const int centroid = find_centroid(find_centroid, root, -1);

    // dist[i] : d(centroid, i)
    vector<int> dist(n);
    auto calc_dist = [&](auto &&self, const int v, const int par, const int d) -> void {
        dist[v] = d;
        for (const int to : tree[v]) {
            if (to == par) {
                continue;
            }
            self(self, to, v, d + 1);
        }
    };
    calc_dist(calc_dist, centroid, -1, 0);

    // out
    i64 ans = 0;
    for (int i = 0; i < n; ++i) {
        ans += dist[i] * c[i];
    }
    cout << ans << endl;
}

Submission Info

Judge Result