dsu on tree学习总结 （树上启发式合并）

模板 ：CF600E Lomsat gelral

题目大意：树的节点有颜色，一种颜色占领了一个子树，当且仅当没有其他颜色在这个子树中出现得比它多。求占领每个子树的所有颜色之和。

最近遇到了一道dsu on tree的题，所以去学了一下，写一下总结。

看到这道题，先考虑暴力：

• 搜索到每个点，暴力搜索一次这个点的子树，时间复杂度是(O(n^2))的。

考虑预处理出每个点的重儿子，每次遍历到一个点，继续遍历出这个点轻儿子的答案，清除这些点对答案的影响。

然后遍历重儿子，同样找出答案，同时保留重儿子对答案的影响。

然后再扫一遍轻儿子，统计出轻儿子的答案。

我们会发现，对于每个点，他的重儿子只会遍历一次，轻儿子的子树会遍历三次。

由重链剖分的知识可得，一个点到根的路径上最多(log_2{n})条轻边，每有一条轻边就会多遍历一次。所以时间是(O(n log_2 n))级别的。

代码

这里给的代码是只走两次轻儿子，用标记来清楚对答案的影响的方法。

#include<cstdio>
#include<cstring>
#include<algorithm>
#define N 200000
#define ll long long
using namespace std;
ll n,x,y,c[N],i,q[N],tot,ans[N],bz[N],son[N],size[N],Max,Sum,bzk[N],edi;
struct edge{
	ll to,next;
}e[N];
void insert(ll x,ll y){
	tot++;
	e[tot].to=y;
	e[tot].next=q[x];
	q[x]=tot;
}
void dfs_pre(ll x,ll father){
	ll i,y;
	size[x]=1;
	for (i=q[x];i;i=e[i].next){
		y=e[i].to;
		if (y!=father){
			dfs_pre(y,x);
			if (size[y]>size[son[x]]) son[x]=y;
			size[x]+=size[y];
		}
	}
}
void get(ll x,ll father,ll son,ll pl,ll id){
	ll i,y;
	if (bzk[c[x]]<pl&&bzk[c[x]]!=id) bzk[c[x]]=id,bz[c[x]]=0;
	bz[c[x]]++;
	if (bz[c[x]]==Max) Sum+=c[x];
	if (bz[c[x]]>Max){
		Max=bz[c[x]];
		Sum=c[x];
	}
	for (i=q[x];i;i=e[i].next){
		y=e[i].to;
		if (y!=father&&y!=son)
			get(y,x,son,pl,id);
	}
}
void dfs(ll x,ll father,ll opt){
	edi++;
	ll i,y,pl=1e9,id=edi;
	for (i=q[x];i;i=e[i].next){
		y=e[i].to;
		if (y!=father&&y!=son[x])
			dfs(y,x,0),Sum=Max=0;
	}
	if (son[x]){
		pl=edi+1;
		dfs(son[x],x,1);
	}
	get(x,father,son[x],pl,id);
	ans[x]=Sum;
}
int main(){
	freopen("dsu_on_tree.in","r",stdin);
	freopen("dsu_on_tree.out","w",stdout);
	scanf("%lld",&n);
	for (i=1;i<=n;i++) scanf("%lld",&c[i]);
	for (i=1;i<n;i++){
		scanf("%lld%lld",&x,&y);
		insert(x,y);
		insert(y,x);
	}
	dfs_pre(1,0);
	dfs(1,0,1);
	Sum=0;
	for (i=1;i<=n;i++) printf("%lld ",ans[i]);
	printf("
");
	return 0;
}