树上数颜色

先考虑暴力怎么做

dfs搜索每个节点，对于当前节点，用迭代遍历其子树并更新答案，在回溯的时候将计数数组清空，时间复杂度为 $O (n^{2})$

注意到，对于每次准备回溯的时候，设当前根节点为 $u$ ，其父亲为 $x$ ，在计算 $x$ 的时候，如果 $u$ 是最后一个dfs计算的（注意此时 $x$ 还没有计算完成，因为我们是dfs遍历树，所以会先计算完 $x$ 的所有儿子），那么从 $u$ 回溯的时候，计数数组就不用清空了，在计算 $x$ 的时候，此时只用迭代遍历 $x$ 的除了 $u$ 的子树并更新答案就好了

那么我们选择的 $u$ 有什么特征呢？一个很自然的想法就是选择子树大小最大的，这样在迭代遍历的时候就遍历得更少了

这是对的，定义重儿子为一个节点的所有儿子中子树大小最大的那个儿子（如果有多个就随便选一个），那么 $u$ 就是重儿子；否则为轻儿子

如果不是很清楚上面的过程，见下面代码

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=1e5+10;
int n,m;
int ans[N];
int c[N];
int Cnt,End[N<<1],Next[N<<1],Last[N];
int cnt[N],sz[N],L[N],R[N],visitime,Node[N],big[N];
int totcol;
void add(int x,int y)
{
	End[++Cnt]=y,Next[Cnt]=Last[x],Last[x]=Cnt;
}
void dfs0(int x,int fa)
{
	L[x]=++visitime;//L表示dfs序的起点 
	Node[visitime]=x;
	sz[x]=1;
	for(int i=Last[x];i;i=Next[i])
	{
		int u=End[i];
		if(u==fa) continue;
		dfs0(u,x);
		sz[x]+=sz[u];
		if(!big[x]||sz[big[x]]<sz[u]) big[x]=u;
	} 
	R[x]=visitime;//R表示子树dfs序的最大值，在dfs序中[L[x],R[x]]就是子树 
}
void add(int u) 
{
  	if(cnt[c[u]]==0) ++totcol;
  	cnt[c[u]]++;
}
void del(int u) 
{
  	cnt[c[u]]--;
  	if(cnt[c[u]]==0) --totcol;
}
void dfs1(int x,int fa,bool keep)//keep表示是否清空计数数组 
{
	for(int i=Last[x];i;i=Next[i])
	{
		int u=End[i];
		if(u==big[x]||u==fa) continue;
		dfs1(u,x,0);
	}
	if(big[x]) dfs1(big[x],x,1);
	for(int i=Last[x];i;i=Next[i])
	{
		int u=End[i];
		if(u==big[x]||u==fa) continue;
		for(int j=L[u];j<=R[u];j++) add(Node[j]);
	}
	add(x);
	ans[x]=totcol;
	if(!keep) 
	for(int i=L[x];i<=R[x];i++) del(Node[i]);
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<n;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		add(u,v),add(v,u);
	}
	for(int i=1;i<=n;i++)
	scanf("%d",&c[i]);
	dfs0(1,0);
	dfs1(1,0,0);
	scanf("%d",&m);
	for(int i=1;i<=m;i++)
	{
		int u;
		scanf("%d",&u);
		printf("%d\n",ans[u]);
	 } 
	return 0;
}

时间复杂度为 $O (n \log n)$

证明：对于每个节点 $u$ ，考虑其对时间复杂度的贡献。显然其会在dfs的时候被遍历一次。除了dfs的这次遍历，其他时候的遍历一定是在计算 $u$ 到根节点的路径上的某个节点 $x$ 的答案的时候，由于 $u$ 没在 $x$ 的重儿子的子树内，导致被迭代遍历到了一次。我们只需要计算这样的 $x$ 有多少个就好了