DSU on tree

DSU on tree

大多数情况下可以替代点分治，并且较为好写
可以理解为一种启发式的暴力，就是树上启发式合并
启发式合并基本思路在于把较小的数据结构向较大的合并，结合重链剖分，可以把暴力做法复杂度压到log级别
例题：给出一棵\(n\)个节点以\(1\)为根的树，节点\(u\)的颜色为\(col_{u}\)，现在对于每个结点\(u\)询问以\(u\)为根的子树里一共出现了多少种不同的颜色。\(n\le 2\times 10^5\)

可以发现，每个节点的答案由其子树和其本身得到，考虑利用这个性质处理问题
为了节省空间，全局开一个桶，桶里维护颜色出现的次数
发现如果切换子树时桶不清空，会造成统计重复
所以先dfs一遍轻儿子，统计轻儿子的答案，消去影响
然后遍历重儿子，得到所有重儿子的信息，保留
接下来再遍历一次轻儿子，得到答案和当前存了重儿子的桶合并，得到整棵树的答案
对于每个点，它到根节点有多少个轻边就被操作多少次，一共n个点，则复杂

具体过程如下：

先预处理出每个节点子树的大小和它的重儿子
\(cnt_{i}\)表示颜色\(i\)的出现次数，\(ans_{u}\)表示结点\(u\)的答案。
遍历节点\(u\)：
先遍历\(u\)的轻（非重）儿子，并计算答案，但不保留遍历后它对\(cnt\)数组的影响；
遍历它的重儿子，保留它对\(cnt\)数组的影响；
再次遍历\(u\)的轻儿子的子树结点，加入这些结点的贡献，以得到\(u\)的答案。

这样，对于一个节点，我们遍历了一次重子树，两次非重子树，显然是最划算的。
通过执行这个过程，我们获得了这个节点所有子树的答案。
为什么不合并第一步和第三步呢？因为\(cnt\)数组不能重复使用，否则空间会太大，需要在\(O(n)\)的空间内完成。
显然若一个节点\(u\)被遍历了\(x\)次，则其重儿子会被遍历\(x\)次，轻儿子（如果有的话）会被遍历\(2*x\)次。
注意除了重儿子，每次遍历完\(cnt\)要清零。

复杂度证明：

对于一棵有\(n\)个节点的树：
根节点到树上任意节点的轻边数不超过\(\log n\)条。我们设根到该节点有\(x\)条轻边该节点的子树大小为\(y\)，显然轻边连接的子节点的子树大小小于父亲的一半（若大于一半就不是轻边了），则\(y<n/2^x\)，显然\(n>2^x\)，所以 \(x<\log n\)。
又因为如果一个节点是其父亲的重儿子，则它的子树必定在它的兄弟之中最多，所以任意节点到根的路径上所有重边连接的父节点在计算答案时必定不会遍历到这个节点，所以一个节点的被遍历的次数等于它到根节点路径上的轻边数\(+1\)（之所以要\(+1\)是因为它本身要被遍历到），所以一个节点的被遍历次数\(=\log n+1\), 总时间复杂度则为\(O(n(\log n+1))=O(n\log n)\)

模板题CF600E代码

#include <iostream>
#define int long long
#define N 100010
using namespace std;
int col[N];
int head[N]; 
int tot;
int ans[N];
int siz[N];
int son[N];
int cnt[N];
int mxn;
int res;
struct edge{
	int to,nxt;
}e[2*N];

void add(int u,int v){
	e[++tot].to=v;
	e[tot].nxt=head[u];
	head[u]=tot;
}

int skp;
//第一遍，树链剖分找重儿子
void init(int u,int fa){
	int maxson=-1;
	siz[u]=1;
	for(int i=head[u];i;i=e[i].nxt){
		int v=e[i].to;
		if(v==fa) continue;
		init(v,u);
		siz[u]+=siz[v];
		if(siz[v]>maxson){
			maxson=siz[v];son[u]=v;
		}
	}
}

//第二遍，计算轻儿子的答案，消除影响
void tra(int u,int fa,int val){
	cnt[col[u]]+=val;
	if(cnt[col[u]]>mxn)mxn=cnt[col[u]],res=col[u];
	else if(cnt[col[u]]==mxn) res+=col[u];
	for(int i=head[u];i;i=e[i].nxt){
		int v=e[i].to;
		if(v==fa||v==skp) continue;
		tra(v,u,val);
	}
}

//第三遍，计算重儿子答案，合并并统计
void dfs(int u,int fa,bool f){
	for(int i=head[u];i;i=e[i].nxt){
		int v=e[i].to;
		if(v==fa||v==son[u]) continue;
		dfs(v,u,0);
	}
	if(son[u]) dfs(son[u],u,1),skp=son[u];
	tra(u,fa,1);skp=0;
	ans[u]=res;
	if(!f) tra(u,fa,-1),res=0,mxn=0;
}
signed main(){
	int n;
	cin>>n;
	for(int i=1;i<=n;i++) cin>>col[i];
	for(int i=1;i<n;i++){
		int u,v;
		cin>>u>>v;
		add(u,v);add(v,u);
	}
	init(1,0);
	dfs(1,0,1);
	for(int i=1;i<=n;i++) cout<<ans[i]<<' ';
	return 0;
}