DSU on tree 学习笔记

DSU on tree 通常用来解决不带修树上子树问题。

主要思想：

1. 剖分。
2. 先搜轻儿子，记录轻儿子子树的答案，删去轻儿子的贡献。
3. 搜重儿子，记录重儿子子树的答案，保留重儿子的贡献。
4. 回溯，重新搜轻儿子，合并轻儿子子树的贡献，构成本子树的答案。

由于每一次合并答案都是把小的子树合并到大的子树里去，因此合并后子树的大小必定大于等于小子树的大小，所以小子树内元素个数至多被合并 $\log n$ 次。共有 $n$ 个元素，所以总合并时间复杂度为 $O(n\log n)$。

CF600E Lomsat gelral

#include<bits/stdc++.h>
using namespace std;

#define ll long long

const int N=1e5+10;

struct edge
{
	int u,v,nxt;
}e[N<<1];
int head[N],tot;
void add(int u,int v)
{
	e[++tot]=edge{u,v,head[u]};
	head[u]=tot;
}

int n;
int col[N];

int siz[N],son[N];
void dfs1(int u,int f)
{
	siz[u]=1;
	for(int i=head[u];i;i=e[i].nxt)
	{
		int v=e[i].v;
		if(v==f) continue;
		dfs1(v,u);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}

int cnt[N];
ll ans[N];

int maxn;
ll sum;

void add(int u,int f,int s)
{
	cnt[col[u]]++;
	if(cnt[col[u]]>maxn) maxn=cnt[col[u]],sum=col[u];
	else if(cnt[col[u]]==maxn) sum+=col[u];
	for(int i=head[u];i;i=e[i].nxt)
	{
		int v=e[i].v;
		if(v==f||v==s) continue;
		add(v,u,s);
	}
}
void sub(int u,int f)
{
	cnt[col[u]]--;
	for(int i=head[u];i;i=e[i].nxt)
	{
		int v=e[i].v;
		if(v==f) continue;
		sub(v,u);
	}
}
void dfs2(int u,int f,int s)
{
	for(int i=head[u];i;i=e[i].nxt)
	{
		int v=e[i].v;
		if(v==f||v==son[u]) continue;
		dfs2(v,u,0);
	}
	if(son[u]) dfs2(son[u],u,1);
	add(u,f,son[u]);
	ans[u]=sum;
	if(!s) sub(u,f),sum=maxn=0;
}
 
int main()
{
	scanf("%d",&n);
	for(int i=1;i<=n;i++) scanf("%d",&col[i]);
	for(int i=1;i<n;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		add(u,v),add(v,u);
	}
	dfs1(1,0),dfs2(1,0,0);
	for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
	return 0;
}