E92 换根DP+倍增 P5666 [CSP-S2019] 树的重心

60分代码（超时，好理解）：

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 3e5 + 10;
int n, rt, s[N], g[N], z[N], p1, p2;
vector<int> G[N];
LL ans, c1[N], c2[N];
void add(LL c[], int x, int k) {for(++x;x<=n+1;x+=x&-x) c[x]+= k;}
LL sum(LL c[], int x) {LL t=0;for(++x;x>=1;x-=x&-x)t+=c[x];return t;}
  
void dfs1(int x, int fa) 
{
    s[x]=1, g[x]= 0;
    for(int y:G[x])if(y!=fa) 
    {
        dfs1(y, x);
        s[x]+= s[y];
        g[x]= max(g[x], s[y]);
    }
    if((n-s[x])*2<=n && g[x]*2<=n) rt=x;
}
void dfs_z(int x, int fa){z[x]=1;for(int y:G[x])if(y!=fa)dfs_z(y, x);}
void dfs3(int x, int fa, int k){add(c1, s[x], k);for(int y:G[x])if(y!=fa)dfs3(y, x, k);}
void dfs2(int x, int fa) 
{
    ans+=rt*( 2*max( s[p1]-(z[x]?s[x]:0), s[p2] ) <= (n-s[x])  ); 
    //判断去掉s[x]后rt是否能成为重心，要考虑x在不在rt的重儿子子树中。  
    if(x!=rt)add(c1, n-s[x], 1);
  
    dfs3(x, fa, -1);
    ans+=x*( sum(c1, n-2*g[x]) - sum(c1, n-2*s[x]-1) );
    dfs3(x, fa, 1);
  
    add(c1, s[x], -1);  
    for(int y:G[x])if(y!=fa) dfs2(y, x);
    add(c1, s[x], 1);
    
    if(x!=rt)add(c1, n-s[x], -1);
}
  
int main() 
{
    int T;scanf("%d", &T);
    while(T--) 
    {
        scanf("%d", &n);
        memset(G, 0, sizeof(G));
        for(int i=1, x, y;i<n;i++) 
        {
            scanf("%d%d", &x, &y);
            G[x].push_back(y);G[y].push_back(x);
        }
        dfs1(1, 0);
        dfs1(rt, 0);
          
        p1=p2=0;s[0]=0;
        for(int x:G[rt])//p1是rt的重儿子，p2是次重儿子 
        {
            if(s[x]>s[p2]) p2=x;
            if(s[p2]>s[p1]) swap(p1, p2);
        }
        memset(z, 0, sizeof(z));dfs_z(p1, rt);//rt的重儿子树所有点的z值都为1 
          
        for(int i=1;i<=n+1;i++) c1[i]=c2[i]=0;
        for(int i=1;i<=n; i++) add(c1, s[i], 1);
        ans = 0;
        dfs2(rt, &0);
        printf("%lld\n", ans);
    }
    return &0;
}

100分代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 3e5 + 10;
int n,rt,s[N],g[N],z[N],p1,p2;
vector<int>G[N];
LL ans,c1[N],c2[N];
void add(LL c[],int x,int k) {for(++x;x<=n+1;x+=x&-x) c[x]+= k;}
LL sum(LL c[],int x) {LL t=0;for(++x;x>=1;x-=x&-x)t+=c[x];return t;}
 
void dfs1(int x,int fa) 
{
    s[x]=1,g[x]= 0;
    for(int y:G[x])if(y!=fa) 
    {
        dfs1(y,x);
        s[x]+= s[y];
        g[x]= max(g[x], s[y]);
    }
    if((n-s[x])*2<=n && g[x]*2<=n) rt=x;
}
void dfs_z(int x,int fa){z[x]=1;for(int y:G[x])if(y!=fa)dfs_z(y,x);}
 
void dfs2(int x,int fa) 
{
    ans+=rt*( 2*max( s[p1]-(z[x]?s[x]:0) , s[p2] ) <= (n-s[x])  ); 
    //判断去掉s[x]后rt是否能成为重心，要考虑x在不在rt的重儿子子树中。 
    if(x!=rt)add(c1, n-s[x], 1);
    
    ans+=x*( sum(c1, n-2*g[x]) - sum(c1, n-2*s[x]-1) );
    ans+=x*( sum(c2, n-2*g[x]) - sum(c2, n-2*s[x]-1) );
 
    add(c2, s[x], 1);
    
    add(c1, s[x], -1);
    for(int y:G[x])if(y!=fa) dfs2(y, x);
    add(c1, s[x], 1);
    
    ans-=x*( sum(c2, n-2*g[x]) - sum(c2, n-2*s[x]-1) );

	if(x!=rt)add(c1, n-s[x], -1);
}
 
int main() 
{
    int T;scanf("%d",&T);
    while(T--) 
    {
        scanf("%d",&n);
        memset(G,0,sizeof(G));
        for(int i=1,x,y;i<n;i++) 
        {
            scanf("%d%d",&x,&y);
            G[x].push_back(y);G[y].push_back(x);
        }
        dfs1(1, 0);
        dfs1(rt, 0);
         
        p1=p2=0;s[0]=0;
        for(int x:G[rt])//p1是rt的重儿子，p2是次重儿子 
        {
            if(s[x]>s[p2]) p2=x;
            if(s[p2]>s[p1]) swap(p1,p2);
        }
        memset(z,0,sizeof(z));dfs_z(p1,rt);//rt的重儿子树所有点的z值都为1 
         
        for(int i=1;i<=n+1;i++) c1[i]=c2[i]=0;
        for(int i=1;i<=n; i++) add(c1, s[i], 1);
        ans = 0;
        dfs2(rt, 0);
        printf("%lld\n",ans);
    }
    return 0;
}