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

视频链接：E92 换根DP+倍增 P5666 [CSP-S2019] 树的重心_哔哩哔哩_bilibili

 

 

 

P5666 [CSP-S2019] 树的重心 - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

// 换根DP+倍增 O(nlogn)
#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;

#define LL long long
#define N 3000005
int read(){
  int x=0,f=1;char c=getchar();
  while(c<'0'||c>'9'){if(c=='-')f=-1;c=getchar();}
  while(c>='0'&&c<='9'){x=x*10+c-'0';c=getchar();}
  return x*f;
}
int T,n;
vector<int> e[N];
int sz[N],f[N],s[N][18],s2[N];
LL ans;

bool check(int u,int n){
  return max(n-sz[u],sz[s[u][0]])<=n/2;
}
void dfs(int u,int fa){
  sz[u]=1,f[u]=fa,s[u][0]=s2[u]=0;
  for(int v:e[u]){
    if(v==fa)continue;
    dfs(v,u);
    sz[u]+=sz[v];
    if(sz[s[u][0]]<sz[v]) s2[u]=s[u][0],s[u][0]=v;
    else if(sz[s2[u]]<sz[v]) s2[u]=v; 
  }
}
void dfs2(int u,int fa){
  int son=s[u][0],szu=sz[u]; f[fa]=u;
  for(int v:e[u]){
    if(v==fa) continue;
    s[u][0]=(v==son?s2[u]:son);
    if(sz[fa]>sz[s[u][0]]) s[u][0]=fa;
    for(int k=1;k<=17;k++)
      s[u][k]=s[s[u][k-1]][k-1];
    sz[u]=n-sz[v],f[u]=0,f[v]=0;
    
    int p=u;
    for(int k=17;k>=0;k--)
      if(sz[u]-sz[s[p][k]]<=sz[u]/2) p=s[p][k];
    ans+=p+f[p]*check(f[p],sz[u]);
    p=v;
    for(int k=17;k>=0;k--)
      if(sz[v]-sz[s[p][k]]<=sz[v]/2) p=s[p][k];
    ans+=p+f[p]*check(f[p],sz[v]);
    dfs2(v,u);
  }
  s[u][0]=son,sz[u]=szu,f[u]=fa; //恢复关系
  for(int k=1;k<=17;k++)
    s[u][k]=s[s[u][k-1]][k-1];
}
int main(){
  T=read();
  while(T--){
    ans=0; n=read();
    for(int i=1;i<=n;i++)e[i].clear();
    for(int i=1;i<n;i++){
      int u=read(),v=read();
      e[u].push_back(v),e[v].push_back(u);
    }
    dfs(1,0);
    // 从i点下跳 1,2,4,8 层的重儿子
    for(int k=1;k<=17;k++)
      for(int i=1;i<=n;i++)
        s[i][k]=s[s[i][k-1]][k-1];
    dfs2(1,0);
    printf("%lld\n",ans);
  }
}

 

// 换根DP+树状数组 O(nlogn)
#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;

#define LL long long
#define N 3000005
LL read(){
  LL x=0,f=1;char c=getchar();
  while(c<'0'||c>'9'){if(c=='-')f=-1;c=getchar();}
  while(c>='0'&&c<='9'){x=x*10+c-'0';c=getchar();}
  return x*f;
}
int T,n,rt;
int sz[N],s[N],s2[N],flag[N];
vector<int> e[N];
LL ans;

struct bit{
  int C[N];
  void init(){memset(C,0,sizeof(C));}
  int lowbit(int u){return u&(-u);}
  void insert(int u,int v){
    u++;
    while(u<=n+1)C[u]+=v,u+=lowbit(u);
  }
  LL query(int u){
    u++;LL res=0;
    while(u)res+=C[u],u-=lowbit(u);
    return res;
  }
}T1,T2;
bool check(int u,int n){
  return max(n-sz[u],sz[s[u]])<=n/2;
}
void dfs(int u,int fa){
  sz[u]=1,s[u]=s2[u]=0;
  for(int v:e[u]){
    if(v==fa)continue;
    dfs(v,u);
    sz[u]+=sz[v];
    if(sz[s[u]]<sz[v])s2[u]=s[u],s[u]=v;
    else if(sz[s2[u]]<sz[v])s2[u]=v;
  }
  if(check(u,n)) rt=u;
}
void dfs2(int u,int fa){
  if(u!=rt){
    flag[u]|=flag[fa];
    T1.insert(sz[fa],-1);
    T1.insert(n-sz[u],1);
    ans+=u*(T1.query(n-2*sz[s[u]])-T1.query(n-2*sz[u]-1));
    ans+=u*(T2.query(n-2*sz[s[u]])-T2.query(n-2*sz[u]-1));
    if(flag[u])ans+=rt*(sz[u]<=n-2*sz[s2[rt]]);
    else ans+=rt*(sz[u]<=n-2*sz[s[rt]]);
  }
  T2.insert(sz[u],1);
  for(int v:e[u]){
    if(v==fa)continue;
    dfs2(v,u);
  }
  if(u!=rt){
    T1.insert(sz[fa],1);
    T1.insert(n-sz[u],-1);
    ans-=u*(T2.query(n-2*sz[s[u]])-T2.query(n-2*sz[u]-1));
  }
}
int main(){
  T=read();
  while(T--){
    ans=rt=0;
    n=read();
    for(int i=1;i<=n;i++)e[i].clear();
    for(int i=1;i<n;i++){
      int u=read(),v=read();
      e[u].push_back(v),e[v].push_back(u);
    }
    dfs(1,0),dfs(rt,0);
    T1.init(),T2.init();
    for(int i=1;i<=n;i++)
      T1.insert(sz[i],1),flag[i]=0;
    flag[s[rt]]=1;
    dfs2(rt,0);
    printf("%lld\n",ans);
  }
}