P5666 - 树的重心 题解

挺 trivial 的一题（虽然当年觉得不可做），由于是 CSP2019 D2T3 就写一波题解罢。

考虑计算每个节点的贡献，也就是要计算每个节点当了几次重心。断开 \((a,fa_a)\) 的话会分成 \(\mathrm{subt}(a)\) 和 \(V-\mathrm{subt}(a)\)​ 两个连通块，分别讨论 \(x\) 属于前者时当重心的次数以及后者。

1. \(x\in\mathrm{subt}(a)\)：此时首先的条件是 \(a\) 是 \(x\) 的假祖先。我们用题面里的定义把「\(x\) 是重心」翻译出来：删除 \(x\) 会分裂出 \(\mathrm{subt}(son)\) 们和 \(\mathrm{subt}(a)-\mathrm{subt}(x)\)。
   1. 对前者写出来就是对所有 \(son\)​​ 有 \(sz_{son}\leq\dfrac{sz_a}2\)​​。即 \(\max\limits_{son}\{sz_{son}\}\leq\dfrac{sz_a}2\)​​。由于我们要数 \(a\)​​，关于 \(x\)​​（以及 \(son\)​​）的都是常数，分离一下得到 \(sz_a\geq 2\max\limits_{son}\{sz_{son}\}\)。
   2. 后者就是 \(sz_a-sz_x\leq\dfrac{sz_a}2\)，即 \(sz_a\leq 2sz_x\)。
2. \(x\in V-\mathrm{subt}(a)\)。此时又要分两种情况：是否有 \(a\in\mathrm{subt}(x)\)，因如果成立的话 \(x\) 的某个子树会被挖掉一块。
   1. \(a\in\mathrm{subt}(x)\)。由于 \(a\neq x\)，设 \(x\) 在 \(a\) 的儿子 \(y\)​ 里。那么要分除了 \(y\) 的儿子树、子树 \(y\)​、\(V-\mathrm{subt}(x)\) 三部分来翻译：
      1. \(\max\limits_{son\neq y}\{sz_{son}\}\leq\dfrac{n-sz_a}2\) 即 \(sz_a\leq n-2\max\limits_{son\neq y}\{sz_{son}\}\)。
      2. \(sz_y-sz_a\leq\dfrac{n-sz_a}2\)​ 即 \(sz_a\geq 2sz_y-n\)​。
      3. \(n-sz_x\leq\dfrac{n-sz_a}2\)​ 即 \(sz_a\leq 2sz_x-n\)​。
   2. \(a\notin\mathrm{subt}(x)\)，此时只要分儿子树和 \(V-\mathrm{subt}(x)-\mathrm{subt}(a)\) 两类：
      1. \(\max\limits_{son}\{sz_{son}\}\leq\dfrac{n-sz_a}2\) 即 \(sz_a\leq n-2\max\limits_{son}\{sz_{son}\}\)。
      2. \(n-sz_a-sz_x\leq\dfrac{n-sz_a}2\) 即 \(sz_a\geq n-2sz_x\)。

上面一共三类，每类都是一些在子树或祖先链上数 \(sz\in[l,r]\) 的数量（对第三种情况就用总的减去祖先后代即可），且总询问数是线性。\(l,r\) 们容易求出，最难的就是 \(\max\limits_{son\neq y}\{sz_{son}\}\)，只要对儿子序列维护一个前后缀最值即可。然后考虑怎么数。子树上的话，dfn 搞下来就是个静态二维数点，主席树即可。由于这是树，也可以 dsu on tree + BIT 2log 或线段树合并 1log，我选择的是主席树。祖先链上的话，不难发现 \(sz\) 是单调的，直接倍增即可。总复杂度线对。

code

#include<bits/stdc++.h>
using namespace std;
#define pb push_back
#define mp make_pair
const int N=300010,LOG_N=20,inf=0x3f3f3f3f;
int n;
vector<int> nei[N];
vector<int> son[N],Mx[N],mX[N];
int fa[N][LOG_N],sz[N],dep[N],dfn[N],mxdfn[N],nowdfn,mng[N];
void dfs(int x=1){
	mng[dfn[x]=mxdfn[x]=++nowdfn]=x;
	sz[x]=1;
	for(int i=1;i<LOG_N;i++)fa[x][i]=fa[fa[x][i-1]][i-1];
	for(int i=0;i<nei[x].size();i++){
		int y=nei[x][i];
		if(y==fa[x][0])continue;
		fa[y][0]=x;
		dep[y]=dep[x]+1;
		dfs(y);
		sz[x]+=sz[y];
		son[x].pb(y),Mx[x].pb(sz[y]),mX[x].pb(sz[y]);
		mxdfn[x]=mxdfn[y];
	}
	for(int i=1;i<Mx[x].size();i++)Mx[x][i]=max(Mx[x][i],Mx[x][i-1]);
	for(int i=int(mX[x].size())-2;i>=0;i--)mX[x][i]=max(mX[x][i],mX[x][i+1]);
	if(Mx[x].empty())Mx[x].pb(-inf);
}
struct segtree{
	int sz,root[N];
	struct node{int lson,rson,cnt;}nd[N<<5];
	#define lson(p) nd[p].lson
	#define rson(p) nd[p].rson
	#define cnt(p) nd[p].cnt
	int nwnd(){return nd[++sz]=node({0,0,0}),sz;}
	void init(){
		sz=root[0]=1;
		nd[0]=nd[1]=node({0,0,0});
	}
	void sprup(int p){cnt(p)=cnt(lson(p))+cnt(rson(p));}
	void add(int x,int v,int p,int np,int tl=1,int tr=n){
		if(tl==tr)return cnt(np)=cnt(p)+1,void();
		int mid=tl+tr>>1;
		if(x<=mid){
			lson(np)=nwnd(),rson(np)=rson(p);
			add(x,v,lson(p),lson(np),tl,mid);
		}
		else{
			rson(np)=nwnd(),lson(np)=lson(p);
			add(x,v,rson(p),rson(np),mid+1,tr);
		}
		sprup(np);
	}
	int _cnt(int l,int r,int p,int tl=1,int tr=n){
		if(l>r)return 0;
		if(l<=tl&&r>=tr)return cnt(p);
		int mid=tl+tr>>1,res=0;
		if(l<=mid)res+=_cnt(l,r,lson(p),tl,mid);
		if(r>mid)res+=_cnt(l,r,rson(p),mid+1,tr);
		return res;
	}
}segt;
int subt(int x,int l,int r){
	l=max(1,l),r=min(n,r);
	return segt._cnt(l,r,segt.root[mxdfn[x]])-segt._cnt(l,r,segt.root[dfn[x]-1]);
}
int ance(int x,int v){
	int cpy=x;
	for(int i=LOG_N-1;~i;i--)if(fa[x][i]&&sz[fa[x][i]]<=v)x=fa[x][i];
	return dep[cpy]-dep[x]+(sz[x]<=v);
}
int ance(int x,int l,int r){
	if(l>r)return 0;
	return ance(x,r)-ance(x,l-1);
}
void mian(){//remember to make it first
	cin>>n;
	nowdfn=0;segt.init();
	for(int i=1;i<=n;i++)nei[i].clear(),son[i].clear(),Mx[i].clear(),mX[i].clear();
	for(int i=1;i<n;i++){
		int x,y;
		scanf("%d%d",&x,&y);
		nei[x].pb(y),nei[y].pb(x);
	}
	dep[1]=1,dfs();
	for(int i=1;i<=n;i++)segt.add(sz[mng[i]],1,segt.root[i-1],segt.root[i]=segt.nwnd());
	long long ans=0;
	for(int i=1;i<=n;i++){
		int tim=0;
		tim+=ance(i,2*Mx[i].back(),2*sz[i])-(2*Mx[i].back()<=n&&n<=2*sz[i]);
//		cout<<tim<<" ";
		for(int j=0;j<son[i].size();j++){
			int y=son[i][j];
			tim+=subt(y,2*sz[y]-n,min(n-2*max(j?Mx[i][j-1]:-inf,j+1<mX[i].size()?mX[i][j+1]:-inf),2*sz[i]-n));
		}
//		cout<<tim<<" ";
		int l=n-2*sz[i],r=n-2*Mx[i].back();
//		cout<<l<<","<<r<<":"<<subt(1,l,r)<<" "<<subt(i,l,r)<<" "<<ance(i,l,r)<<"!!\n";
		tim+=subt(1,l,r)-subt(i,l,r)-ance(i,l,r)+(l<=sz[i]&&sz[i]<=r);
//		cout<<tim<<"!!!\n";
		ans+=1ll*i*tim;
	}
	cout<<ans<<"\n";
}
int main(){
	int testnum;
	cin>>testnum;
	while(testnum--)mian();
	return 0;
}