【模闆】樹上啟發式合并

模闆題

#include<bits/stdc++.h>
using namespace std;
#define N 100005
int n,a[N],ans[N],son[N],siz[N];
vector<int>g[N];
int f[N],sum,m;
inline void dfs1(int x,int fa)
{
	siz[x]=1;
	int ml=0,md=0;
	for(int i=0;i<g[x].size();i++)
		if(g[x][i]!=fa)
		{
			dfs1(g[x][i],x);
			siz[x]+=siz[g[x][i]];
			if(siz[g[x][i]]>ml)
				ml=siz[g[x][i]],md=g[x][i];
		}
	son[x]=md;
}
inline void dfs2(int x,int fa,int d,int s1)
{
	f[a[x]]+=d;
	if(d==1&&f[a[x]]==1)
		sum++;
	if(f[a[x]]==0)
		sum--;
	for(int i=0;i<g[x].size();i++)
		if(g[x][i]!=fa&&g[x][i]!=s1)
			dfs2(g[x][i],x,d,s1);
}
inline void dfs(int x,int fa,bool f1)
{
	int l1=g[x].size();
	for(int i=0;i<l1;i++)
		if(g[x][i]!=fa&&g[x][i]!=son[x])
			dfs(g[x][i],x,0);
	if(son[x])
		dfs(son[x],x,1);
	dfs2(x,fa,1,son[x]);
	ans[x]=sum;
	if(!f1)
		dfs2(x,fa,-1,0);
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<n;i++)
	{
		int x,y;
		scanf("%d%d",&x,&y);
		g[x].push_back(y);
		g[y].push_back(x);
	}
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	dfs1(1,0);
	dfs(1,0,1);
	scanf("%d",&m);
	for(int i=1;i<=m;i++)
	{
		int x;
		scanf("%d",&x);
		printf("%d\n",ans[x]);
	}
}