U41492 数颜色 树上启发式合并

树上启发式合并的经典问题。
首先离线处理答案。用一个颜色数组记录某种颜色是否存在过，然后由于不同子树统计答案的时候互不影响，我们枚举子树的时候要清空之前的数组，但是这就带来了一个问题，我们无法向上合并答案了。
启发式合并的做法是我们先考虑树剖预处理，然后对于统计一个结点的答案，先暴力递归所有轻儿子的答案不保留数组，然后递归重儿子并且保留数组，然后再次递归所有的轻儿子这次也保留数组，从而向上更新答案。
由于轻边的个数不超过$O(log(n))$，因此每个结点最多向上暴力合并$log$次，而重链的个数也不会超过$O(log(n))$个，因此复杂度不会超过$O(logn)$。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int inf=0x3f3f3f3f;
const ll INF=LONG_LONG_MAX;
const int N=1e5+7;
vector<int> G[N];
int sz[N],son[N],res[N];
bool check[N];
int c[N]; 
void dfs1(int u,int fa) {
	sz[u]=1;
	for(auto &v:G[u]) {
		if(v==fa) continue;
		dfs1(v,u);
		sz[u]+=sz[v];
		if(sz[v]>sz[son[u]]) 
			son[u]=v; 
	}
} 
int dfs2(int u,int fa,int keep) {   
	if(keep) {
		for(auto &v:G[u]) {
			if(v==fa||v==son[u]) continue;
				dfs2(v,u,keep); 
		}
	}
	int ans=0;
	if(son[u])  ans+=dfs2(son[u],u,keep);
	for(auto &v:G[u]) {
		if(v==fa||v==son[u]) continue;
		ans+=dfs2(v,u,0); 
	}
	if(!check[c[u]]) {
		ans++;
		check[c[u]]=1;
	}
	// keep 1 轻边不保留check  
	// keep 0 保留check 
	if(keep) {
		res[u]=ans;
		if(u!=son[fa]) 
			memset(check,0,sizeof(check));	
	}
	return ans;
} 
int main() {
	int n;
	scanf("%d",&n);
	for(int i=1;i<n;i++) {
		int u,v;
		scanf("%d%d",&u,&v);
		G[u].push_back(v);
		G[v].push_back(u);
	}
	for(int i=1;i<=n;i++)
		scanf("%d",&c[i]);
	dfs1(1,0); 
	dfs2(1,0,1);
	int q;
	scanf("%d",&q);
	while(q--) {
		int x;
		scanf("%d",&x);
		printf("%d\n",res[x]);
	}
	return 0;
}