计蒜客 蒜头君的建设方案 （树形dp）

 

#include <bits/stdc++.h>

using namespace std;

#define inf 0x3f3f3f3f

const int maxn=1e6+10;

typedef long long ll;

ll f[maxn];
ll dp[maxn];
ll size[maxn];

int cnt;
int head[maxn*2];
int n;

ll ans;

struct node{
    int to;
    int next;
    int len;
}p[maxn*2];

void addedge(int u,int v,int w){
    p[++cnt].to=v;
    p[cnt].len=w;
    p[cnt].next=head[u];
    head[u]=cnt;
}

inline void init(){
    cnt=0;
    ans=0;
    memset(head,-1,sizeof(head));
}

void get_deep(int u,int fa){
    f[u]=size[u]=1;
    for(int e=head[u];e!=-1;e=p[e].next){
        if(p[e].to!=fa){
            get_deep(p[e].to,u);
            size[u]+=size[p[e].to];
            f[u]+=f[p[e].to]+size[p[e].to];
        }
    }
}

void get_ans(int u,int fa){
    for(int e=head[u];e!=-1;e=p[e].next){
        if(p[e].to!=fa){
            dp[p[e].to]=dp[u]+n-2*size[p[e].to];
            ans+=abs(dp[u]-dp[p[e].to])*p[e].len;
            get_ans(p[e].to,u);
        }
    }
}

int main()
{
    init();
    int u,v,w;
    scanf("%d",&n);
    for(int i=1;i<n;i++){
        scanf("%d%d%d",&u,&v,&w);
        addedge(u,v,w);
        addedge(v,u,w);
    }
    get_deep(1,0);
    dp[1]=f[1];
    get_ans(1,0);
    printf("%lld\n",ans);
    return 0;
}