树上数颜色

先考虑暴力怎么做

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$有多少个就好了

定义轻/重边表示连接父节点与轻/重儿子的边，显然满足条件的$x$的个数就是从$u$到根节点的路径上轻边的个数，下面证明轻边的个数不超过$\log n$

那么这个与启发式合并有什么关系呢？其实就是证明过程换一种理解方式。我们从$u$往上走，每经过一条轻边，子树大小就会至少扩大一倍，就跟启发式合并一样了，最多扩大$\log n$次