[换根DP][倍增]luogu P5666 树的重心

题面

https://www.luogu.com.cn/problem/P5666

分析

对于一棵以i为根的树来说，它的重心必然在其size大于等于sumsize/2的子树中。

那么断掉一条边e(u,v)时，我们对于断掉边的u,v进行讨论，然后向他们的重儿子倍增直到满足其size≤sumsize/2。

具体实现时可能存在两个重心，所以要判断一下找到的点的重儿子和其父亲。

然后换根的时候维护一下size和father就行了。

代码

#include <iostream>
#include <cstdio>
#include <cstring>
using namespace std;
typedef long long ll;
const int N=3e5+10;
struct Graph {
    int v,nx;
}g[2*N];
int cnt,list[N];
int t,n,sz[N],h[2][N],bg[N],d[N][20],f[2][N];
ll ans;

void Add(int u,int v) {g[++cnt]=(Graph){v,list[u]};list[u]=cnt;}

void DFS1(int u) {
    sz[u]=1;h[0][u]=h[1][u]=bg[u]=0;
    for (int i=list[u];i;i=g[i].nx)
        if (g[i].v!=f[0][u]) {
            f[1][g[i].v]=f[0][g[i].v]=u;DFS1(g[i].v);sz[u]+=sz[g[i].v];
            if (sz[g[i].v]>sz[h[0][u]]) h[1][u]=h[0][u],bg[u]=h[0][u]=g[i].v;
            else if (sz[g[i].v]>sz[h[1][u]]) h[1][u]=g[i].v;
        }
    d[u][0]=h[0][u];
    for (int i=1;i<=18;i++) d[u][i]=d[d[u][i-1]][i-1];
}

void DFS2(int u) {
    for (int i=list[u],x;i;i=g[i].nx)
        if(g[i].v!=f[0][u]) {
            sz[u]=n-sz[g[i].v];
            if (g[i].v==h[0][u]) bg[u]=h[1][u]; else bg[u]=h[0][u];
            if (sz[bg[u]]<sz[f[0][u]]) bg[u]=f[0][u];
            f[1][u]=f[1][g[i].v]=0;d[u][0]=bg[u];
            for (int j=1;j<=18;j++) d[u][j]=d[d[u][j-1]][j-1];
            x=u;
            for (int j=18;j>=0;j--) if (sz[d[x][j]]>sz[u]/2) x=d[x][j];
            if (max(sz[bg[x]],sz[u]-sz[x])<=sz[u]/2) ans+=x;
            if (max(sz[bg[bg[x]]],sz[u]-sz[bg[x]])<=sz[u]/2) ans+=bg[x];
            if (max(sz[bg[f[1][x]]],sz[u]-sz[f[1][x]])<=sz[u]/2) ans+=f[1][x];
            x=g[i].v;
            for (int j=18;j>=0;j--) if (sz[d[x][j]]>sz[g[i].v]/2) x=d[x][j];
            if (max(sz[bg[x]],sz[g[i].v]-sz[x])<=sz[g[i].v]/2) ans+=x;
            if (max(sz[bg[bg[x]]],sz[g[i].v]-sz[bg[x]])<=sz[g[i].v]/2) ans+=bg[x];
            if (max(sz[bg[f[1][x]]],sz[g[i].v]-sz[f[1][x]])<=sz[g[i].v]/2) ans+=f[1][x];
            f[1][u]=g[i].v;
            DFS2(g[i].v);
        }
    sz[u]=n-sz[f[1][u]=f[0][u]];d[u][0]=bg[u]=h[0][u];
    for (int i=1;i<=18;i++) d[u][i]=d[d[u][i-1]][i-1];
}

int main() {
    for (scanf("%d",&t);t;t--) {
        scanf("%d",&n);memset(list,cnt=0,sizeof list);
        for (int i=1,u,v;i<n;i++) scanf("%d%d",&u,&v),Add(u,v),Add(v,u);
        ans=0;DFS1(1);DFS2(1);
        printf("%lld\n",ans);
    }
}