[hdu6567]Cotree

Time Limit: 2000/1000 MS (Java/Others)
Memory Limit: 262144/262144 K (Java/Others)

分數：2500
Problem Description
Avin has two trees which are not connected. He asks you to add an edge between them to make them connected while minimizing the function $∑^n_{i=1}∑^n_{j=i+1}dis(i,j)$, where $dis(i,j)$ represents the number of edges of the path from $i$ to $j$. He is happy with only the function value.

Input
The first line contains a number $n(2<=n<=100000)$. In each of the following $n−2$ lines, there are two numbers $u$ and $v$, meaning that there is an edge between $u$ and $v$. The input is guaranteed to contain exactly two trees.

Output
Just print the minimum function value.

Sample Input

3
1 2

Sample Output

4

題意：
給定一棵樹的$n-2$條邊，讓你加上一條邊，讓他變爲一棵樹，並且滿足$∑^n_{i=1}∑^n_{j=i+1}dis(i,j)$最小，$dis(i,j)$是$i$到$j$的距離。
題解：
首先我們對兩個樹進行dp，對每個節點計算$g[x]$，表示爲以$x$爲根節點
的子樹中所有節點到根節點的距離之和。
然後我們再dfs一次計算對於一個節點來說，同一個樹的其他所有點到這個點的距離之和$p[i]$
那麼$(\sum p[i])/2$是不加邊之前的答案。那麼加邊之後的答案，只需要找到兩個數各自最小的兩個$p[i]$,設爲$a1$,$a2$,答案就是$(a1+c1)*c2+a2*c1$

#include<bits/stdc++.h>
#define ll long long
using namespace std;
int n;
vector<int>G[200004];
int bl[200004],sz[200004];
ll g[200004],ans[4],res;
int c[4],rt[4];
void dfs(int x,int bt){
    bl[x]=bt;
    ++c[bt];
    sz[x]=1;g[x]=0;
    for(int i=0;i<G[x].size();i++){
        int to=G[x][i];
        if(bl[to])continue;
        dfs(to,bt);
        sz[x]+=sz[to];
        g[x]+=sz[to];
        g[x]+=g[to];
    }
}
void dfs2(int x,int fa,ll val){
    res+=val+g[x];
    //cout<<x<<" "<<val+g[x]<<" "<<bl[x]<<" || "<<endl;
    if(ans[bl[x]]==-1||(ans[bl[x]]!=-1&&ans[bl[x]]>val+g[x])){
        ans[bl[x]]=val+g[x];
    }
    for(int i=0;i<G[x].size();i++){
        int to=G[x][i];
        if(to==fa)continue;
        dfs2(to,x,val+c[bl[x]]-2*sz[to]+g[x]-g[to]);
    }
}
int w33ha(){
    res=0;
    for(int i=1;i<=n;i++){
        G[i].clear();
        g[i]=0;
        ans[i]=0;
        bl[i]=0;
    }
    c[1]=0;c[2]=0;
    rt[1]=0;rt[2]=0;
    ans[1]=-1;ans[2]=-1;
    for(int i=1;i<=n-2;i++){
        int u,v;scanf("%d%d",&u,&v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    int blo=0;
    for(int i=1;i<=n;i++){
        if(!bl[i]){
            dfs(i,++blo);
            rt[blo]=i;
        }
    }
    //for(int i=1;i<=n;i++)cout<<g[i]<<" ";cout<<endl;
    for(int i=1;i<=blo;i++){
        dfs2(rt[i],0,0);
    }
    //for(int i=1;i<=n;i++)cout<<bl[i]<<" ";cout<<endl;
    res>>=1;
    //cout<<res<<endl;
    //cout<<" "<<ans[1]<<" "<<ans[2]<<endl;
    printf("%lld\n",res+(ans[1]+c[1])*(c[2])+(ans[2]*c[1]));
    return 0;
}
int main(){
    while(scanf("%d",&n)!=EOF)w33ha();
    return 0;
}