树上启发式合并（dsu on tree）

dsu on tree

$\text{dsu}$一般指$\text{disjoint set union}$，即并查集。$\text{dsu on tree}$指树上合并与查询操作，但它的实现和普通的并查集并无关联，两者的共同点仅仅在于都能合并集合和查询而已。

$\text{dsu on tree}$，可以称为树上启发式合并，是一种巧妙的暴力。用一个全局数组存储结果，对于每棵子树，有以下操作：

• 先遍历轻儿子，处理完轻儿子后将数组清零（不能用$memset$，要再遍历一次来清零）
• 遍历重儿子，遍历完不用清零，再遍历，将轻儿子合并到重儿子上去，其合并结果存储于全局数组
• 用此时的全局数组来计算父亲

重儿子： 对于一个非叶节点$u$，设$v$是$u$的儿子，且以$v$为根的子树包含的节点比以$u$的其他儿子为根的子树包含的节点都多，则称$v$为$u$的重儿子

轻儿子： 对于一个非叶节点$u$，在$u$的各个儿子中，除了重儿子，都是轻儿子

重边： 重儿子连向父亲的边

轻边： 轻儿子连向父亲的边

重链： 若干条重边连接而成的链

运用树上启发式合并，可以将时间复杂度降至$O(n\log n)$。

为什么这样做时间复杂度就能降为$O(n\log n)$呢？

根据轻重链划分的思想，任意一条从某个节点到根的路径上轻边的数量不超过$\log n$条，而重链是被轻边分隔的，所以数量也不超过$\log n$条。每棵子树到父亲的边为轻边才需要做一次$pt$操作，所以最多做$\log n$次。每次$pt$操作可以看作是轻儿子向重儿子的合并操作。每个节点最多合并$\log n$次，所以时间复杂度为$O(n\log n)$。

例题

传送门

给一棵根为1的树，每次询问子树颜色种类数。

先看一下暴力怎么做。

code

void pt(int u,int fa,int fl){
	cnt[a[u]]+=fl;
	if(fl==1&&cnt[a[u]]==1) ++now;
	if(fl==-1&&cnt[a[u]]==0) --now;
	for(int i=r[u];i;i=l[i]){
		if(d[i]==fa) continue;
		pt(d[i],u,fl);
	}
}
void dfs(int u,int fa){
	for(int i=r[u];i;i=l[i]){
		if(d[i]==fa) continue;
		dfs(d[i],u);
	}
	pt(u,fa,1);
	ans[u]=now;
	pt(u,fa,-1);
}

时间复杂度为$O(n^2)$。

然后来看看树上启发式合并的代码。

code

void dfs1(int u,int fa){
	siz[u]=1;
	for(int i=r[u];i;i=l[i]){
		if(d[i]==fa) continue;
		dfs1(d[i],u);
		siz[u]+=siz[d[i]];
		if(siz[d[i]]>siz[son[u]]) son[u]=d[i];
	}
}
void pt(int u,int fa,int fl){
	cnt[a[u]]+=fl;
	if(fl==1&&cnt[a[u]]==1) ++now;
	if(fl==-1&&cnt[a[u]]==0) --now;
	for(int i=r[u];i;i=l[i]){
		if(d[i]==fa) continue;
		pt(d[i],u,fl);
	}
}
void dfs2(int u,int fa){
	for(int i=r[u];i;i=l[i]){
		if(d[i]==fa||d[i]==son[u]) continue;
		dfs2(d[i],u);
		pt(d[i],u,-1);
	}
	if(son[u]) dfs2(son[u],u);
	++cnt[a[u]];
	if(cnt[a[u]]==1) ++now;
	for(int i=r[u];i;i=l[i]){
		if(d[i]!=fa&&d[i]!=son[u]) pt(d[i],u,1);
	}
	ans[u]=now;
}

其中dfs1是用来求重儿子的，dfs2就用来求答案。在遍历各节点的过程中，我们只对轻儿子作清零操作，重儿子求完再遍历各个轻儿子来求父亲的答案。这样就能将时间复杂度降为$O(n\log n)$了。

完整代码如下。

code

#include<bits/stdc++.h>
using namespace std;
int n,m,x,y,tot=0,now=0,a[100005],d[200005],l[200005],r[200005],siz[100005],son[100005],cnt[100005],ans[100005];
void add(int xx,int yy){
	l[++tot]=r[xx];d[tot]=yy;r[xx]=tot;
}
void dfs1(int u,int fa){
	siz[u]=1;
	for(int i=r[u];i;i=l[i]){
		if(d[i]==fa) continue;
		dfs1(d[i],u);
		siz[u]+=siz[d[i]];
		if(siz[d[i]]>siz[son[u]]) son[u]=d[i];
	}
}
void pt(int u,int fa,int fl){
	cnt[a[u]]+=fl;
	if(fl==1&&cnt[a[u]]==1) ++now;
	if(fl==-1&&cnt[a[u]]==0) --now;
	for(int i=r[u];i;i=l[i]){
		if(d[i]==fa) continue;
		pt(d[i],u,fl);
	}
}
void dfs2(int u,int fa){
	for(int i=r[u];i;i=l[i]){
		if(d[i]==fa||d[i]==son[u]) continue;
		dfs2(d[i],u);
		pt(d[i],u,-1);
	}
	if(son[u]) dfs2(son[u],u);
	++cnt[a[u]];
	if(cnt[a[u]]==1) ++now;
	for(int i=r[u];i;i=l[i]){
		if(d[i]!=fa&&d[i]!=son[u]) pt(d[i],u,1);
	}
	ans[u]=now;
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<n;i++){
		scanf("%d%d",&x,&y);
		add(x,y);add(y,x);
	}
	for(int i=1;i<=n;i++){
		scanf("%d",&a[i]);
	}
	dfs1(1,0);
	dfs2(1,0);
	scanf("%d",&m);
	while(m--){
		scanf("%d",&x);
		printf("%d\n",ans[x]);
	}
	return 0;
}