[USACO10MAR] 伟大的奶牛聚集 - 树形dp

Posted 2021-03-12 mollnn

每个点有重数，求到所有点距离最小的点

就是魔改的重心了

#include <bits/stdc++.h>
using namespace std;

#define int long long

const int N = 1000005;

vector <pair<int,int> > g[N];
int siz[N],f[N],vis[N],sum[N],c[N],n,m,t1,t2,t3,tot;

void dfs1(int p) {
    vis[p]=1;
    siz[p]=c[p];
    for(int i=0;i<g[p].size();i++) {
        int q=g[p][i].first, w=g[p][i].second;
        if(vis[q]==0) {
            dfs1(q);
            sum[p]+=sum[q]+siz[q]*w;
            siz[p]+=siz[q];
        }
    }
}

void dfs2(int p) {
    vis[p]=1;
    for(int i=0;i<g[p].size();i++) {
        int q=g[p][i].first, w=g[p][i].second;
        if(vis[q]==0) {
            f[q] = f[p] - siz[q]*w + (tot-siz[q])*w;
            dfs2(q);
        }
    }
}

signed main() {
    ios::sync_with_stdio(false);
    cin>>n;
    for(int i=1;i<=n;i++) cin>>c[i], tot+=c[i];
    for(int i=1;i<n;i++) {
        cin>>t1>>t2>>t3;
        g[t1].push_back(make_pair(t2,t3));
        g[t2].push_back(make_pair(t1,t3));
    }
    dfs1(1);
    memset(vis,0,sizeof vis);
    f[1]=sum[1];
    dfs2(1);
    int ans = 1e+18;
    for(int i=1;i<=n;i++) {
        ans = min(ans, f[i]);
    }
    cout<<ans<<endl;
}