[CSP-S2019] 树的重心

Description

给定一棵树，对于每一条边，求出割掉该条边后，两棵树的所有可能的重心，输出重心编号和。

Solution

一个性质是，如果当前点不是重心，那么重心只可能是往父亲方向走或者重儿子方向走。转换一下，如果割掉 \(u\to v\) 这条边，分别建两棵以 \(u\) 和 \(v\) 为根的树，那么找重心只需要走重儿子。所以需要处理的只有一个换根操作。我们可以先预处理一个以 \(1\) 为根的所有节点的倍增数组。考虑当前删掉 \(u\to v\)，重构树，那么实际上被改变的祖父关系只有 \(u\) 的根缀上的点，考虑增量更新，回溯的时候再还原即可。详见代码。

#include<stdio.h>
#include<algorithm>
using namespace std;

const int N=3e5+7;

inline int read(){
    int x=0,flag=1; char c=getchar();
    while(c<'0'||c>'9'){if(c=='-')flag=0;c=getchar();}
    while(c>='0'&&c<='9'){x=(x<<1)+(x<<3)+c-48;c=getchar();}
    return flag? x:-x;
}

struct E{
    int next,to;
}e[N<<1];

int head[N],cnt=0;
inline void add(int id,int to){
    e[++cnt]=(E){head[id],to};
    head[id]=cnt;
    e[++cnt]=(E){head[to],id};
    head[to]=cnt;
}

int son[N][2],f[N][19],fa[N],sz[N],S[N],Son[N];

void dfs(int u){
    sz[u]=1,son[u][0]=son[u][1]=0;
    for(int i=head[u];i;i=e[i].next){
        int v=e[i].to;
        if(v==fa[u]) continue;
        fa[v]=u,dfs(v);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u][0]])
            son[u][1]=son[u][0],son[u][0]=v;
        else if(sz[v]>sz[son[u][1]]) son[u][1]=v;
    }
    f[u][0]=son[u][0];
    for(int i=1;i<19;i++) f[u][i]=f[f[u][i-1]][i-1];
}

inline int check(int u,int sz){
    return max(S[Son[u]],sz-S[u])<=sz/2? /*printf("%d ",u),*/u:0; 
}

long long ans=0;
void Dfs(int u,int Fa){
    for(int i=head[u];i;i=e[i].next){
        int v=e[i].to;
        if(Fa==v) continue;
        S[u]=sz[1]-sz[v]; fa[u]=fa[v]=0;
        //split it into Two subtrees,regard u,v as root
        Son[u]=(son[u][0]!=v? son[u][0]:son[u][1]);
        if(S[Fa]>S[Son[u]]) Son[u]=Fa;
        f[u][0]=Son[u]; int rt=u;
        for(int j=1;j<19;j++) f[u][j]=f[f[u][j-1]][j-1];
        for(int j=18;~j;j--)
            if(f[rt][j]&&S[u]-S[f[rt][j]]<=S[u]/2) rt=f[rt][j];
    //    printf("[%d,%d] ",u,v);
        ans+=check(rt,S[u])+check(fa[rt],S[u])+check(Son[rt],S[u]);
        rt=v; for(int j=18;~j;j--)
            if(f[rt][j]&&S[v]-S[f[rt][j]]<=S[v]/2) rt=f[rt][j];
        ans+=check(rt,S[v])+check(fa[rt],S[v])+check(Son[rt],S[v]);// printf("\n");
        fa[u]=v,Dfs(v,u);
    }
    Son[u]=f[u][0]=son[u][0],fa[u]=Fa,S[u]=sz[u];
    for(int i=1;i<19;i++) f[u][i]=f[f[u][i-1]][i-1];
}

int main(){
    int T=read();
    while(T--){
        int n=read(); cnt=0;
        for(int i=1;i<=n;i++) head[i]=fa[i]=0;
        for(int i=2;i<=n;i++) add(read(),read()); dfs(1);
        for(int i=1;i<=n;i++)
            S[i]=sz[i],Son[i]=son[i][0];
        ans=0,Dfs(1,0); printf("%lld\n",ans);
    }
}