2019CCPC-江西省赛 -A Cotree (树形DP,求树上一点到其他点的距离之和)

我是傻逼我是傻逼

#include<bits/stdc++.h>
using namespace std;
const int maxn=4e5+50;
typedef long long ll;
const ll mod=998244353;

int Laxt[maxn],Next[maxn],To[maxn],cnt;
void add(int u,int v){
    Next[++cnt]=Laxt[u];
    Laxt[u]=cnt;
    To[cnt]=v;
    Next[++cnt]=Laxt[v];
    Laxt[v]=cnt;
    To[cnt]=u;
}
ll dp[maxn],sz[maxn];
int vis[maxn];
void dfs(int u,int p){
    sz[u]=1;vis[u]=1;
    for(int i=Laxt[u];i;i=Next[i]){
        int v=To[i];
        if(v==p)continue;
        dfs(v,u);
        sz[u]+=sz[v];
        dp[u]+=dp[v]+sz[v];
    }
}
ll mi[5],id[5];
void dfs1(int u,int pre,int o){
    if(mi[o]>dp[u]){
        mi[o]=dp[u];id[o]=u;
    }
   // cout<<"u="<<u<<" "<<"dp[u]="<<dp[u]<<endl;
    for(int i=Laxt[u];i;i=Next[i]){
        int v=To[i];
        if(v==pre)continue;

        dp[u]-=dp[v]+sz[v];
        sz[u]-=sz[v];
        dp[v]+=dp[u]+sz[u];
        sz[v]+=sz[u];
        dfs1(v,u,o);

        sz[v]-=sz[u];
        dp[v]-=dp[u]+sz[u];
        sz[u]+=sz[v];
        dp[u]+=dp[v]+sz[v];
    }
}
ll sum=0;
void dfs2(int u,int pre){
    sum+=dp[u];
   // cout<<"u="<<u<<" "<<dp[u]<<endl;
    for(int i=Laxt[u];i;i=Next[i]){
        int v=To[i];
        if(v==pre)continue;
        dp[u]-=dp[v]+sz[v];
        sz[u]-=sz[v];
        dp[v]+=dp[u]+sz[u];
        sz[v]+=sz[u];
        dfs2(v,u);
        sz[v]-=sz[u];
        dp[v]-=dp[u]+sz[u];
        sz[u]+=sz[v];
        dp[u]+=dp[v]+sz[v];
    }
}

int main(){
    std::ios::sync_with_stdio(false);
    int n;
    cin>>n;
    for(int i=1;i<=n-2;i++){
        int u,v;
        cin>>u>>v;
        add(u,v);
    }
    int tot=0;
    for(int i=0;i<4;i++)mi[i]=1e18;
    for(int i=1;i<=n;i++){
        if(!vis[i]){
            dfs(i,0);
           /* for(int i=1;i<=n;i++){
                cout<<"i="<<i<<" "<<dp[i]<<" "<<sz[i]<<endl;
            }*/
            dfs1(i,0,++tot);
        }
    }
    add(id[1],id[2]);
    memset(dp,0,sizeof(dp));
    memset(sz,0,sizeof(sz));
    dfs(1,0);
    dfs2(1,0);
    cout<<sum/2<<endl;
    return 0;
}