E85 换根DP P2986 [USACO10MAR] Great Cow Gathering G

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 1e5+10;
vector<pair<int, int>> G[N];
ll cnt, siz[N], c[N], f[N], d[N], ans=(1LL<<60);

void dfs(int x, int xfa)
{
    siz[x] = c[x];
    for (auto i : G[x]) if (i.first != xfa)
    {
        int y = i.first, w = i.second;
        dfs(y, x);
        siz[x] += siz[y];
        f[x] = f[x] + f[y] + siz[y] * w;
    }
}

void dp(int x, int xfa)
{
    for (auto i : G[x]) if (i.first != xfa)
    {
        int y = i.first, w = i.second;
        d[y] = d[x] - siz[y] * w + (cnt - siz[y]) * w;
        dp(y, x);
    }
}

int main()
{
    int n; scanf("%d", &n);
    cnt = 0; for (int i = 1; i <= n; i++) scanf("%lld", &c[i]), cnt += c[i];
    for (int i = 1, x, y, w; i < n; i++)
    {
        scanf("%d%d%d", &x, &y, &w);
        G[x].push_back(make_pair(y, w));
        G[y].push_back(make_pair(x, w));
    }
    dfs(1, 0);
    d[1] = f[1]; dp(1, 0);
    ll ans = d[1]; for (int i = 2; i <= n; i++) ans = min(ans, d[i]);
    printf("%lld\n", ans);
    return 0;
}