【点分治】Cotree HDU

题目

链接：http://acm.hdu.edu.cn/showproblem.php?pid=6567

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 ∑ni=1∑nj=i+1dis(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

 

 

代码

#include<iostream>
#include<algorithm>
#include<cstring>
#include<cstdio>
#include<vector>
#include<map>
#include<stack>
using namespace std;
typedef long long ll;
const double eps=1e-8;
const ll N=1e5+5;
const int mod=1e9+7;

int vis[N],size[N];

vector<int> e[N];

//重心
int Gans=1e9;
int pos=0;
ll ans=0;
void Gdfs(int x,int n){
    size[x]=1;
    int mp=0;

    for(int i=0;i<e[x].size();i++){
        int y=e[x][i];
        if(vis[y])continue;
        vis[y]=1;
        Gdfs(y,n);
        size[x]+=size[y];
        mp=max(mp,size[y]);
        vis[y]=0;
    }
    mp=max(mp,n-size[x]);
    if(mp<Gans){
        Gans=mp;
        pos=x;
    }
}

//树节点计数
int Num=0;
void Ndfs(int x){
    Num++;
    for(int i=0;i<e[x].size();i++){
        int y=e[x][i];
        if(vis[y])continue;
        vis[y]=1;
        Ndfs(y);
        vis[y]=0;
    }
}

//染色，计数
int Rvis[N];
void Rdfs(int x){
    Num++;
    for(int i=0;i<e[x].size();i++){
        int y=e[x][i];
        if(vis[y]||Rvis[y])continue;
        Rvis[y]=1;
        Rdfs(y);

    }
}

//预处理子树值
ll Ssuml[N], Snuml[N];
ll Snum,Ssd, Ssum;
ll SSsuml[N],SSnuml[N];
void dfs(int x){
    Ssum+=Ssd;
    Snum++;
    for(int i=0;i<e[x].size();i++){
        int y=e[x][i];
        if(vis[y])continue;
        vis[y]=1;
        Ssd++;
        dfs(y);
        vis[y]=0;
        Ssd--;

    }

}

//点分治
void dfz(int x){
    //换根为重心
    Num=0;
    vis[x]=1;
    Ndfs(x);
    vis[x]=0;
    int n=Num;
    Gans=1e9;
    Gdfs(x,n);
    x=pos;
    vis[x]=1;
    int sn=0;
    
    //求子树值，来计算经过根的路径长
    for(int i=0;i<e[x].size();i++){
        int y=e[x][i];
        if(vis[y])continue;
        vis[y]=1;
        Ssum=Snum=0;
        Ssd=1;
        dfs(y);
        vis[y]=0;
        Ssuml[sn]=Ssum;
        Snuml[sn++]=Snum;
    }

    for(int i=0;i<=sn;i++){
        SSnuml[i]=0;
        SSsuml[i]=0;
    }
    for(int i=sn-1;i>=0;i--){
        SSnuml[i]=SSnuml[i+1]+Snuml[i];
        SSsuml[i]=SSsuml[i+1]+Ssuml[i];
    }

    if(sn)ans+=SSsuml[0];
    for(int i=0;i<sn;i++){
        ans+=SSnuml[i+1]*Ssuml[i]+SSsuml[i+1]*Snuml[i];
    }

    //分治
    for(int i=0;i<e[x].size();i++){
        int y=e[x][i];
        if(vis[y])continue;
        dfz(y);
    }



}

signed main(){
    
    int n;
    scanf("%d",&n);
    int a,b;
    for(int i=0;i<n-2;i++){
        scanf("%d%d",&a,&b);
        e[a].emplace_back(b);
        e[b].emplace_back(a);
    }

    int r1=1,r2;
    int r1s,r2s;
    //找第一颗树
    Num=0;
    Rdfs(1);
    r1s=Num;
    //找第二颗树
    for(int i=2;i<=n;i++){
        if(Rvis[i])continue;
        Num=0;
        Rdfs(i);
        r2s=Num;
        r2=i;
    }
    //找第一颗树重心
    Gans=1e9;
    Gdfs(r1,r1s);
    int pos1=pos;
    //找第二颗树重心
    Gans=1e9;
    Gdfs(r2,r2s);
    int pos2=pos;
    //连接重心
    e[pos1].emplace_back(pos2);
    e[pos2].emplace_back(pos1);
    //点分治
    dfz(1);

    printf("%lld
",ans);

    scanf(" ");
}