[Codeforces 600E]Lomsat gelral(dsu on tree)

Sample Input

4
1 2 3 4
1 2
2 3
2 4

Sample Output

10 9 3 4

题意
给出一个n个节点的树，以1为树根，每个点有一种颜色。一种颜色占领了一个子树当且仅当没有其他颜色在这个子树中出现次数比它多（允许次数相同）。求占领每个子树的所有颜色之和。

思路
这题看了题解才会的
dsu on tree（树上并查集？但是感觉跟并查集没什么关系，反而有树剖的思想）
算是小技巧
在统计一棵子树时，先找出重儿子，统计答案的时候优先进入所有轻儿子，之后清除对应计算信息，再进入重儿子，目的是保留重儿子所在子树的信息。向上合并时直接将所有轻儿子的信息加入重儿子的信息。

（复杂度证明也是看题解）
考虑树剖，从根到某一点的路径上，不超过$O(logN)$条轻边，所以任何一个点最多被作为轻儿子计算$O(logN)$次，总时间复杂度$O(NlogN)$

AC代码

#include<bits/stdc++.h>
#include<math.h>
using namespace std;
typedef long long ll;
const int maxn=1e5+5;
int n,ma;
int h[maxn],p[maxn*2],nxt[maxn*2],c[maxn];
long long ans[maxn],siz[maxn],son[maxn],cnt[maxn],su[maxn];//cnt:对颜色的计数 su:数量为i的所有颜色之和 
void add(int now,int f,int val)
{
	for(int i=h[now];i;i=nxt[i])
	{
		if(p[i]!=f&&!son[p[i]]) add(p[i],now,val);//不计算被标记为重儿子的子树（因为已经计算过了） 
	}
	su[cnt[c[now]]]-=c[now];
	cnt[c[now]]+=val;
	su[cnt[c[now]]]+=c[now];
	if(val+1)
	{
		if(su[ma+1]) ma++;
	}else
	{
		if(!su[ma])ma--;
	}
	return;
}
int sea(int now,int f)//预处理，计算每个子树的大小 
{
	int mx,son_;
	siz[now]=1;
	for(int i=h[now];i;i=nxt[i])
	if(p[i]!=f)
	{
		sea(p[i],now);
		siz[now]+=siz[p[i]];
	}
	return 0;	
}
int dfs(int now,int f,int keep)//keep标记用来说明是否要保留影响 
{
	int son_,mx;
	mx=son_=-1;
	for(int i=h[now];i;i=nxt[i])
	if(p[i]!=f)if(siz[p[i]]>mx)mx=siz[p[i]],son_=p[i];//记录重儿子 
	for(int i=h[now];i;i=nxt[i])
	if(p[i]!=f&&p[i]!=son_)
	{
		dfs(p[i],now,0);//计算所有非重儿子的子树，清除影响 
	}
	if(son_!=-1)
	{son[son_]=1;dfs(son_,now,1);}//标记重儿子，计算重儿子子树，不清除影响 
	add(now,f,1);//这里加入所有非重儿子的子树，在这个函数里检查标记看是否为重儿子，是则不计算对应子树 
	if(son_!=-1)son[son_]=0;//清除标记，为后面的清理做准备（如果keep为1也要清） 
	ans[now]=su[ma];
	if(!keep)//不保留影响则清除，此时重儿子标记已经删除，所以会清理整棵子树的影响 
	add(now,f,-1),ma=0;
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<=n;i++)scanf("%d",&c[i]);
	for(int i=1;i<n;i++)
	{
		int x,y;
		scanf("%d%d",&x,&y);
		p[i*2-1]=y;nxt[i*2-1]=h[x];h[x]=i*2-1;
		p[i*2]=x;nxt[i*2]=h[y];h[y]=i*2;
	}
	sea(1,0);
	dfs(1,0,1);
	for(int i=1;i<=n;i++)printf("%lld ",ans[i]);
	return 0;
}