[luogu5666]树的重心
考虑枚举一个点k，求其为重心的方案数
暴力的做法是，将其作为根搜索，设最大子树大小为s1，次大为s2，对割掉的子树分类讨论：
1.在子树中，分两种情况（都可以用线段树合并来做）
(1)从s1中切掉一棵大小为s3的子树，应该满足$2max(s2,s1-s3)\le n-s3$，即$2s1-n\le s3\le n-2s2$
(2)从其他子树中切掉一棵大小为s3的子树，应该满足$2s1\le n-s3$，即$s3\le n-2s1$
2.是父亲，那么割掉的边再分为两类（这些东西也需要再根据父亲是不是最大子树来讨论）
(1)割掉的边是直接到根的路径，那么割掉的子树大小s3就是n-割出来的子树，可以再搜一遍不断的维护当前节点到根的路径上所有子树大小，用权值线段树来维护区间和
(2)割掉的边是其他边，直接对最终线段树合并到根的线段树上查询即可（注意这样会错误计算第(1)种情况，要注意在第一个中抵消掉，即再对子树大小打上-1标记）
1 #include<bits/stdc++.h> 2 using namespace std; 3 #define N 300005 4 #define mid (l+r>>1) 5 struct ji{ 6 int nex,to; 7 }edge[N<<1]; 8 struct node{ 9 int sum,ls,rs; 10 }f[N*50]; 11 int E,V,t,n,x,y,head[N],r[N],fi[N],se[N],sz[N]; 12 long long ans; 13 void up(int k){ 14 f[k].sum=f[f[k].ls].sum+f[f[k].rs].sum; 15 } 16 void update(int &k,int l,int r,int x,int y){ 17 if (!k)k=++V; 18 if (l==r){ 19 f[k].sum+=y; 20 return; 21 } 22 if (x<=mid)update(f[k].ls,l,mid,x,y); 23 else update(f[k].rs,mid+1,r,x,y); 24 up(k); 25 } 26 int query(int k,int l,int r,int x,int y){ 27 if ((!k)||(l>y)||(x>r))return 0; 28 if ((x<=l)&&(r<=y))return f[k].sum; 29 return query(f[k].ls,l,mid,x,y)+query(f[k].rs,mid+1,r,x,y); 30 } 31 int merge(int k1,int k2){ 32 if ((!k1)||(!k2))return k1+k2; 33 if ((!f[k1].ls)&&(!f[k1].rs)){ 34 f[k1].sum+=f[k2].sum; 35 return k1; 36 } 37 f[k1].ls=merge(f[k1].ls,f[k2].ls); 38 f[k1].rs=merge(f[k1].rs,f[k2].rs); 39 up(k1); 40 return k1; 41 } 42 void add(int x,int y){ 43 edge[E].nex=head[x]; 44 edge[E].to=y; 45 head[x]=E++; 46 } 47 void dfs(int k,int fa){ 48 fi[k]=sz[k]=0; 49 sz[k]=1; 50 for(int i=head[k];i!=-1;i=edge[i].nex) 51 if (edge[i].to!=fa){ 52 dfs(edge[i].to,k); 53 sz[k]+=sz[edge[i].to]; 54 if (sz[edge[i].to]<fi[k])se[k]=max(se[k],sz[edge[i].to]); 55 else{ 56 se[k]=fi[k]; 57 fi[k]=sz[edge[i].to]; 58 } 59 } 60 if (n-sz[k]<fi[k])se[k]=max(se[k],n-sz[k]); 61 else{ 62 se[k]=fi[k]; 63 fi[k]=n-sz[k]; 64 } 65 for(int i=head[k];i!=-1;i=edge[i].nex) 66 if (edge[i].to!=fa){ 67 if (sz[edge[i].to]!=fi[k])ans+=1LL*k*query(r[edge[i].to],1,n,1,n-2*fi[k]); 68 else ans+=1LL*k*query(r[edge[i].to],1,n,max(2*fi[k]-n,1),n-2*se[k]); 69 r[k]=merge(r[k],r[edge[i].to]); 70 } 71 if (fi[k]!=n-sz[k])ans-=1LL*k*query(r[k],1,n,1,n-2*fi[k]); 72 else ans-=1LL*k*query(r[k],1,n,max(2*fi[k]-n,1),n-2*se[k]); 73 update(r[k],1,n,sz[k],1); 74 } 75 void dfs2(int k,int fa){ 76 update(r[1],1,n,sz[k],-1); 77 if (k>1)update(r[1],1,n,n-sz[k],1); 78 if (fi[k]!=n-sz[k])ans+=1LL*k*query(r[1],1,n,1,n-2*fi[k]); 79 else ans+=1LL*k*query(r[1],1,n,max(2*fi[k]-n,1),n-2*se[k]); 80 for(int i=head[k];i!=-1;i=edge[i].nex) 81 if (edge[i].to!=fa)dfs2(edge[i].to,k); 82 update(r[1],1,n,sz[k],1); 83 update(r[1],1,n,n-sz[k],-1); 84 } 85 int main(){ 86 scanf("%d",&t); 87 while (t--){ 88 E=V=ans=0; 89 memset(r,0,sizeof(r)); 90 memset(f,0,sizeof(f)); 91 memset(head,-1,sizeof(head)); 92 scanf("%d",&n); 93 for(int i=1;i<n;i++){ 94 scanf("%d%d",&x,&y); 95 add(x,y); 96 add(y,x); 97 } 98 dfs(1,0); 99 dfs2(1,0); 100 printf("%lld\n",ans); 101 } 102 }
