CF 600E Lomsat gelral （dsu on tree）

题目链接：https://codeforces.com/contest/600/problem/E

题意：给一个n的点的树，根为1。子树u的支配颜色为出现次数最多的颜色（可能有多个），求每个子树中支配颜色的和。

思路：常规套路，其实只需要维护颜色的个数。其实计算每个子树答案的时候，算一下最大次数以及和就行。由于这是做的第一个题，所以我另外维护了一个ma数组（出现的最大次数）、sum数组（当前出现次数最多的颜色的和）。

代码：

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 1e5+10;
vector<int> g[N];
int n,c[N],sz[N],num[N],ma[N];
ll sum[N],ans[N];
void getsz(int u,int fa){
	int siz=g[u].size();
	sz[u]=1;
	for(int i=0;i<siz;i++){
		int v=g[u][i];
		if(v==fa) continue;
		getsz(v,u);
		sz[u]+=sz[v];
	}
}
bool big[N];
void dfs1(int u,int fa,int x){//计算轻儿子的贡献 
	sum[num[c[u]]]-=c[u];
	num[c[u]]+=x;
	ma[u]=num[c[u]];	
	sum[num[c[u]]]+=c[u];		
	int siz=g[u].size();
	for(int i=0;i<siz;i++){
		int v=g[u][i];
		if(v==fa||big[v]) continue ;
		dfs1(v,u,x);
		ma[u]=max(ma[v],ma[u]);
	}	
}
void dfs(int u,int fa,bool keep){
	int mx=-1,bigc=-1;
	int siz=g[u].size();
	//找重儿子 
	for(int i=0;i<siz;i++){
		int v=g[u][i];
		if(v==fa) continue;
		if(sz[v]>mx) bigc=v,mx=sz[v];
	}
	//计算所有轻儿子的答案，然后清除	
	for(int i=0;i<siz;i++){
		int v=g[u][i];
		if(v==fa||v==bigc) continue;
		dfs(v,u,0);
	}
	//计算重儿子的答案，不清除 
	if(bigc!=-1){
		dfs(bigc,u,1);
		big[bigc]=1;
	}
	//计算父亲和轻儿子的贡献	
	dfs1(u,fa,1);
	//计算重儿子的贡献 
	if(bigc!=-1) ma[u]=max(ma[u],ma[bigc]);
	//得到父亲的答案 
	ans[u]=sum[ma[u]]; 
	if(bigc!=-1) big[bigc]=0;
	//如果不是重儿子，删除自己子树的贡献 
	if(keep==0)	
		dfs1(u,fa,-1);	
}
int main(void){
	scanf("%d",&n);
	for(int i=1;i<=n;i++) scanf("%d",&c[i]);
	for(int i=1;i<n;i++){
		int u,v;
		scanf("%d%d",&u,&v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	getsz(1,0);
	dfs(1,0,0);
	for(int i=1;i<=n;i++)
		printf("%lld%c",ans[i],i==n?'\n':' ');
	return 0;	
}