[CSP-S2019] 树的重心

#include<cstdio>
#include<cmath>
#include<cstring>
#include<algorithm>
#define maxn 300005
#define maxp 20
#define mem(a) memset(a,0,sizeof(a))
#define ll long long 
using namespace std;

int T,n,i,j,k,x,y;
int em,e[maxn*2],nx[maxn*2],ls[maxn];
int fa[maxn],sz[maxn],gs[maxn];
int Sz[maxn],Gs[maxn],f[maxn][maxp],F[maxn][maxp];
ll ans;

void read(int &x){
    x=0; char ch=getchar();
    for(;ch<'0'||ch>'9';ch=getchar());
    for(;ch>='0'&&ch<='9';ch=getchar()) x=x*10+ch-'0';
}

void insert(int x,int y){
    em++; e[em]=y; nx[em]=ls[x]; ls[x]=em;
    em++; e[em]=x; nx[em]=ls[y]; ls[y]=em;
}

void DFS(int x,int p){
    sz[x]=1,fa[x]=p;
    for(int i=ls[x];i;i=nx[i]) if (e[i]!=p){
        DFS(e[i],x),sz[x]+=sz[e[i]];
        if (!gs[x]||sz[gs[x]]<sz[e[i]])     
            gs[x]=e[i];
    }
    f[x][0]=gs[x],Gs[x]=gs[x];
    for(int i=1;i<maxp;i++) f[x][i]=f[f[x][i-1]][i-1];
}

void upd(int x){
    for(int i=1;i<maxp;i++)     
        F[x][i]=F[F[x][i-1]][i-1];
}

void Get(int st){
    int N=Sz[st];
    int y1=st,y2=st,c=0;
    for(int i=maxp-1;i>=0;i--) if (F[y1][i]&&Sz[F[y1][i]]>=(N+1)/2)
        y1=F[y1][i],c+=1<<i;
    if (c){
        c--;
        for(int i=maxp-1;i>=0;i--) if (c>=(1<<i))
            c-=1<<i,y2=F[y2][i];
        int mx=min(max(N-Sz[y1],Sz[Gs[y1]]),max(N-Sz[y2],Sz[Gs[y2]]));
        if (max(N-Sz[y1],Sz[Gs[y1]])==mx) ans+=y1;
        if (max(N-Sz[y2],Sz[Gs[y2]])==mx) ans+=y2;
    } else ans+=y1;
}

int totT,To[maxn],pre[maxn],nex[maxn];
void DFS2(int x,int p){
    if (x!=1) Get(x);
    totT=0;
    for(int i=ls[x];i;i=nx[i]) if (e[i]!=p) To[++totT]=e[i];
    int mx=To[1];
    for(int i=2;i<=totT;i++){
        pre[To[i]]=mx;
        if (Sz[To[i]]>Sz[mx]) mx=To[i];
    }
    mx=To[totT];
    for(int i=totT-1;i>=1;i--){
        nex[To[i]]=mx;
        if (Sz[To[i]]>Sz[mx]) mx=To[i];
    }
    
    for(int i=ls[x];i;i=nx[i]) if (e[i]!=p){
        Gs[x]=fa[x];
        if (pre[e[i]]&&Sz[pre[e[i]]]>Sz[Gs[x]]) Gs[x]=pre[e[i]];
        if (nex[e[i]]&&Sz[nex[e[i]]]>Sz[Gs[x]]) Gs[x]=nex[e[i]];
        F[x][0]=Gs[x],upd(x);
        Sz[x]=n-Sz[e[i]];
        Get(x);
        DFS2(e[i],x);
    }
    Sz[x]=sz[x],Gs[x]=gs[x];
    memcpy(F[x],f[x],sizeof(f[x]));
}

int main(){
    freopen("centroid.in","r",stdin);
    freopen("centroid.out","w",stdout);
    read(T);
    while (T--){
        read(n),ans=0;
        em=0,mem(ls),mem(f),mem(sz),mem(gs),mem(pre),mem(nex);
        for(i=1;i<n;i++) read(x),read(y),insert(x,y);
        DFS(1,0);
        memcpy(Sz,sz,sizeof(sz));
        memcpy(F,f,sizeof(F));
        DFS2(1,0);
        printf("%lld\n",ans);
    }
}