树上启发式合并（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 delif(mod)add(cur);elsedel(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，因此该树是颜色平衡树。