CF600E Lomsat gelral（dsu on tree）

蒟蒻在CF留下足迹
集训中Day5的一道题
典型的dsu on tree
递归处理轻儿子，计算轻儿子答案，
然后消去轻儿子对于答案的影响。
然后递归处理重儿子，不消去影响，最后加入所有轻儿子贡献，计算答案。
还可以将树处理成序列，用莫队进行处理
dsu on tree时间复杂度O（nlogn）
莫队时间复杂度O（n√n）
显然前者更优
PS 需要开long long，亲测最后一个数据点int无法通过

#include <bits/stdc++.h>
using namespace std;
#define maxn 100005
#define ll long long
ll cnt=0,top=0,n;
ll head[maxn],size[maxn],deep[maxn],f[maxn],son[maxn],sum[maxn],num[maxn],c[maxn],ans[maxn];
bool vis[maxn];
struct node
{
	int to,next;
}e[maxn*2];
inline ll read()
{
	ll w=1,s=0; char ch=getchar();
	while(ch<'0' || ch>'9'){if(ch=='-')w=-1; ch=getchar();}
	while(ch>='0' && ch<='9'){s=s*10+ch-'0'; ch=getchar();}
	return w*s;
}
void add(int x,int y)
{
	cnt++;
	e[cnt].to=y;
	e[cnt].next=head[x];
	head[x]=cnt;
}
void dfs1(int u,int fa)
{
	size[u]=1; deep[u]=deep[fa]+1; f[u]=fa;
	for(int i=head[u];i;i=e[i].next)
	{
		int v=e[i].to;
		if(v==fa) continue;
		dfs1(v,u);
		size[u]+=size[v];
		if(size[v]>size[son[u]]) son[u]=v;
	}
}
void update(int u,int fa,int tp)
{
	sum[num[c[u]]]-=c[u];
	num[c[u]]+=tp;
	sum[num[c[u]]]+=c[u];
	if(sum[top+1]) top++;
	if(!sum[top]) top--;
	for(int i=head[u];i;i=e[i].next)
	{
		int v=e[i].to;
		if(v==fa || vis[v]) continue;
		update(v,u,tp);
	}
}
void dfs(int u,int fa,int tp)
{
	
	for(int i=head[u];i;i=e[i].next)
	{
		int v=e[i].to;
		if(v==fa || v==son[u]) continue;
		dfs(v,u,0);
	}
	if(son[u]) 
	{
		dfs(son[u],u,1); vis[son[u]]=true;
	}
	update(u,fa,1); vis[son[u]]=false;
	ans[u]=sum[top];
	if(!tp) update(u,fa,-1);
}
int main()
{
	n=read();
	for(int i=1;i<=n;i++) c[i]=read();
	for(int i=1;i<n;i++)
	{
		int a=read(),b=read();
		add(a,b); add(b,a);
	}
	dfs1(1,0);
	dfs(1,0,1);
	for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
	return 0;
}