[CSU 1811]Tree Intersection(dsu on tree)

Descroption
Bobo has a tree with $n$ vertices numbered by 1,2,…,n and (n-1) edges. The $i$-th vertex has color $c_i$, and the i-th edge connects vertices $a_i$ and $b_i$.
Let $C(x,y)$ denotes the set of colors in subtree rooted at vertex $x$ deleting edge $(x,y)$.
Bobo would like to know $R_i$ which is the size of intersection of $C(a_i,b_i)$ and $C(b_i,a_i)$ for all $1\le i\le (n-1)$. (i.e. |$C(a_i,b_i)\cap C(b_i,a_i)$|)
Input
The input contains at most 15 sets. For each set:
The first line contains an integer $n(2\le n\le 10^5).$
The second line contains $n$ integers $c_1$,$c_2$,…,$c_n$ $(1\le c_i\le n)$.
The $i$-th of the last (n-1) lines contains 2 integers $a_i,b_i(1\le a_i,b_i\le n)$.
Output
For each set, (n-1) integers $R_1,R_2,\dots ,R_{n-1}$.
Sample Input

4
1 2 2 1
1 2
2 3
3 4
5
1 1 2 1 2
1 3
2 3
3 5
4 5

Sample Output

1
2
1
1
1
2
1

题意
_多组数据
给一棵$n$个点的树，每个点上有一种颜色，共n-1条边，第i条边连接结点$a_i,b_i$。
定义$C(x,y)$是去掉边（x，y）之后以x为根的子树拥有的颜色集合。
定义$R_i$ 是$C(a_i,b_i)$ 与 $C(b_i,a_i)$交集的大小。
顺序输出$R_1,R_2,\dots \dots ,R_{n-1}$

思路
实质是计算去掉一条边后形成的两棵子树拥有的颜色数的交集的大小。题目没有指定树根，不妨以1为根。那么问题转化为：枚举每条边，计算去掉这条边之后得到的不以1为根的子树与以1为根的子树拥有的颜色集合的交集大小。如果事先统计了所有颜色的数量，那么对于每种情况，只需要统计根不为1的子树每个点的颜色数量(cnt[]数组)就可以反推出另一棵子树的颜色数量(cc[]数组)，进而推得集合。
然后就是dsu on tree，每次扫完一棵子树后向上合并就好了。
添加子树时cc[i]+1,则cnt[i]-1，然后比较一下cc[i]与cnt[i]改变前后的值，来决定是否更新交集大小。
注意到一个结点可能有多个子结点，但只有一个（或者没有）父结点，所以如果有一条边（f[i],i），则我们将这条边的答案存入ans[i]而不是ans[f[i]]。输出时也需要判断（因为是枚举每条边输出），这里我使用了深度来代替f[]判断父结点。

多组数据记得清空数组
AC代码

#include<bits/stdc++.h>
#include<math.h>
using namespace std;
typedef long long ll;
const int maxn=1e5+5;
int n,ma;
int h[maxn],p[maxn*2],nxt[maxn*2],c[maxn];
long long ans[maxn],siz[maxn],son[maxn],cnt[maxn],cc[maxn],dep[maxn];//cnt:对颜色的计数 su:数量为i的所有颜色之和 
void add(int now,int f,int val)
{
	for(int i=h[now];i;i=nxt[i])
	{
		if(p[i]!=f&&!son[p[i]]) add(p[i],now,val);//不计算被标记为重儿子的子树（因为已经计算过了） 
	}
	int temp1=cc[c[now]]&&cnt[c[now]];
	cc[c[now]]-=val;
	cnt[c[now]]+=val;
	int temp2=cc[c[now]]&&cnt[c[now]];
	ma+=temp2-temp1;
	return;
}
int sea(int now,int f)//预处理，计算每个子树的大小 
{
	int mx,son_;
	siz[now]=1;dep[now]=dep[f]+1;
	for(int i=h[now];i;i=nxt[i])
	if(p[i]!=f)
	{
		sea(p[i],now);
		siz[now]+=siz[p[i]];
	}
	return 0;	
}
int dfs(int now,int f,int keep)//keep标记用来说明是否要保留影响 
{
	int son_,mx;
	mx=son_=-1;
	for(int i=h[now];i;i=nxt[i])
	if(p[i]!=f)if(siz[p[i]]>mx)mx=siz[p[i]],son_=p[i];//记录重儿子 
	for(int i=h[now];i;i=nxt[i])
	if(p[i]!=f&&p[i]!=son_)
	{
		dfs(p[i],now,0);//计算所有非重儿子的子树，清除影响 
	}
	if(son_!=-1)
	{son[son_]=1;dfs(son_,now,1);}//标记重儿子，计算重儿子子树，不清除影响 
	add(now,f,1);//这里加入所有非重儿子的子树，在这个函数里检查标记看是否为重儿子，是则不计算对应子树 
	if(son_!=-1)son[son_]=0;//清除标记，为后面的清理做准备（如果keep为1也要清） 
	if(son_==-1)ans[now]=(cc[c[now]]!=0);
	else ans[now]=ma;
	if(!keep)//不保留影响则清除，此时重儿子标记已经删除，所以会清理整棵子树的影响 
	add(now,f,-1),ma=0;
}
int main()
{
	while(scanf("%d",&n)!=EOF)
	{
		
		for(int i=1;i<=n;i++)
		{	
			scanf("%d",&c[i]);
			cc[c[i]]++;
		}
		for(int i=1;i<n;i++)
		{
			int x,y;
			scanf("%d%d",&x,&y);
			p[i*2-1]=y;nxt[i*2-1]=h[x];h[x]=i*2-1;
			p[i*2]=x;nxt[i*2]=h[y];h[y]=i*2;
		}
		sea(1,0);
		dfs(1,0,1);
		for(int i=1;i<n;i++)printf("%lld\n",dep[p[i*2-1]]>dep[p[i*2]]?ans[p[i*2-1]]:ans[p[i*2]]);
		for(int i=1;i<=n;i++)cnt[c[i]]--,h[i]=0;
	}
	return 0;
}