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;

}