树上启发式合并（dsu on tree）学习

声明：本文部分内容摘自OI Wiki网站。详情可自行查看学习。

洛谷 P9233

题目实际上是蓝桥杯 2023 年 A 组省赛的一道题。题干大致的意思是，给定一个含有 $n$ 个结点，并且以 $1$ 为根的一棵树，每个节点 $i$ 都有一个颜色 $C_i$ ，问这棵树有多少个子树是颜色平衡树。其中，如果一棵树中存在的每种颜色的结点个数都相同，则我们称它是一棵颜色平衡树。
输入：一行一个 $n$ ，后面的 $n$ 行每行一个数对 $C_i,F_i)$ 分别表示结点 $i$ 的颜色和父亲结点。（保证 $F_1=0$ ）。
输出：这棵树有多少棵子树是颜色平衡树。

题目分析

感觉题目最大的难点是， $C_i\leq 200000$ 。因为我们最朴素的思想就是：看一棵树是否是颜色平衡树，需要将其子树的颜色情况统计起来。比如，开一个数组c[1..N]，其中子树中颜色为i的结点有c[i]个。通过统计一棵树所有子树的c可以得到这棵树的c。
但是现在有一个问题，如果按照上面的思路，显然有 $N\geq C_i$ ；如果每碰到一个子树就开一个c[N]，肯定会MLE，并且由于每次合并子树的颜色情况时需要遍历这个数组c，多半也会TLE。
于是我们又产生了一个想法，用一个map来代替c[N]。因为不是每棵子树都会出现所有的颜色，我们开map不会为没有出现的颜色分配空间，显然空间使用就大大减少了。这种情况下合并子树的颜色情况时，只需要遍历所有子树的map，一一加到树上就行了。
经过实践，使用map的方法确实不会MLE，但是它TLE了（70 pts）。

接触到了树上启发式合并这一思想之后，迅速写代码，于是AC。树上启发式合并的核心思想就是保留重儿子。在下面的代码中有注释说明。

AC代码

#include<iostream>

using namespace std;

int n,c[200005],ccnt[200005],cccnt,sz[200005],ncolor[200005],ecnt=1,fedge[200005],ledge[200005],ans,bigchild[200005];
struct {
	int end,next;
}edge[200005];

void buildarc(int begin,int end){
	if(!begin)
		return;
	if(!fedge[begin])
		fedge[begin]=ledge[begin]=ecnt;
	else{
		edge[ledge[begin]].next=ecnt;
		ledge[begin]=ecnt;
	}
	edge[ecnt++].end=end;
}

inline void addcc(int cc){
	if(!ccnt[cc])
		cccnt++;
	ccnt[cc]++;
}
inline void delcc(int cc){
	ccnt[cc]--;
	if(!ccnt[cc])
		cccnt--;
}

void add(int node){
	if(c[ncolor[node]])
		delcc(c[ncolor[node]]);
	addcc(c[ncolor[node]]+1);
	c[ncolor[node]]++;
}
void del(int node){
	c[ncolor[node]]--;
	delcc(c[ncolor[node]]+1);
	if(c[ncolor[node]])
		addcc(c[ncolor[node]]);
}

inline int isbalanced(){return cccnt==1?1:0;}

// dfs0 用来求每个子树的大小的，也就是这个子树有多少个结点。
int dfs0(int cur,int father){
	int childsize,maxchild=0;
	sz[cur]=1;
	for(int e=fedge[cur];e;e=edge[e].next){
		if(edge[e].end==father)
			continue;
		sz[cur]+=(childsize=dfs0(edge[e].end,cur));
		if(childsize>maxchild){
			maxchild=childsize;
			bigchild[cur]=edge[e].end;
		}
	}
	return sz[cur];
}

// 这个是用来加入/删除某一棵树的，mod == 1 是加入，mod == 0 是删除。
void dfs2(int cur,int father,bool mod){//1 add 0 del
	if(mod)
		add(cur);
	else
		del(cur);
	for(int e=fedge[cur];e;e=edge[e].next)
		if(edge[e].end!=father)
			dfs2(edge[e].end,cur,mod);
}
/***
 dfs1 是树上启发式合并的核心算法。一个数有多个子树，其中最大的子树的树根成为这个树的重儿子，其它的儿子是轻儿子。树上启发式算法在每个子树上的操作分为 3 步：
 1.遍历所有轻儿子子树，查看情况，不保留轻儿子子树对原树 c[N] 数组的影响。
 2.查看重儿子子树的情况，并且保留该子树对原树 c[N] 数组的影响。
 3.重新加入所有轻儿子子树，从而得到原树的情况。

@param cur 当前结点
@param father 当前结点的父节点，用 father = 0 表示没有父节点
@param remain 是否要保留 cur 为根的子树对其父树的影响。也即 cur 是否是重儿子。
***/
void dfs1(int cur,int father,bool remain){
	//遍历所有轻儿子
	for(int e=fedge[cur];e;e=edge[e].next)
		if(edge[e].end!=father && edge[e].end!= bigchild[cur])
			dfs1(edge[e].end,cur,false);
	//查看重儿子
	if(bigchild[cur])
		dfs1(bigchild[cur],cur,true);
	//加入根节点
	add(cur);
	//重新加入所有轻儿子子树
	for(int e=fedge[cur];e;e=edge[e].next)
		if(edge[e].end!=father && edge[e].end!= bigchild[cur])
			dfs2(edge[e].end,cur,true);
	ans+=isbalanced();
	//如果要求不保留影响，cur 为根的树的所有贡献
	if(!remain)
		dfs2(cur,father,false);
}

int main(){
	int f;
	cin>>n;
	for(int i=1;i<=n;i++){
		scanf("%d%d",&ncolor[i],&f);
		buildarc(f,i);
	}
	return dfs0(1,0),dfs1(1,0,true),cout<<ans,0;
}

AC代码中保留了c[N]数组，其含义与上文所述相同，即颜色的出现次数。可以看到代码中还出现了ccnt数组和cccnt变量。这是一个比较精妙的处理方式，可以在 $O (1)$ 的时间复杂度内判断一棵树是否是颜色平衡树。ccnt[i] = j表示当前共有j个颜色的c值是i，即颜色出现次数的出现次数。cccnt是颜色出现次数的出现次数的出现次数，当且仅当cccnt == 1的时候，该树是一颗颜色平衡树。
上面的描述可能有些烧脑，可以用题目给的样例来说明。样例输入：

读者可以根据输入的含义自行画图。对于根节点1而言，颜色1,2,3分别出现了2,2,2次，所以c[1,2,3] = {2,2,2}。根据定义，有3个颜色的出现次数是2，所以ccnt[2] = 3。而对于其它的i != 2，都有cccnt[i] = 0，所以cccnt = 1，因此该树是颜色平衡树。