CF600E Lomsat gelral （树上启发式合并模板）

题目链接：点此跳转

题目大意：一棵树有n个结点，每个结点都是一种颜色，每个颜色有一个编号，求树中每个子树的最多的颜色编号的和。

解题思路：这道是树启的模板题，我们先跑一遍dfs处理出轻重儿子，然后每次保留重儿子的贡献，消除轻儿子的贡献，这样时间复杂度可以降到log级别。

Code：

#include<iostream>
#include<cstring>
#include<queue>
using namespace std;
typedef long long ll;
const int maxn=1e5+7;

struct edge{
	int v,nxt;
}e[maxn<<1];

int h[maxn],cntt;

int col[maxn],cnt[maxn],sz[maxn],w[maxn],n;
ll ans[maxn],mc,cc;

void init(){
	memset(h,-1,sizeof h);
	cntt=0;
}

void add(int u,int v){
	e[cntt].v=v;
	e[cntt].nxt=h[u];
	h[u]=cntt++;
}

void dfs1(int u,int fa){            //处理轻重儿子
	sz[u]=1;
	for(int i=h[u];~i;i=e[i].nxt){
		int v=e[i].v;
		if(v==fa) continue;
		dfs1(v,u);
		
		if(sz[v]>sz[w[u]]) w[u]=v;
		sz[u]+=sz[v];
	}
}

void calc(int u,int fa,int o,int W){                //计算结果
	cnt[col[u]]+=o;
	if(cnt[col[u]]==mc) cc+=col[u];
	else if(cnt[col[u]]>mc) mc=cnt[col[u]],cc=col[u];
	
	for(int i=h[u];~i;i=e[i].nxt){
		int v=e[i].v;
		if(v==fa||v==W) continue;
		calc(v,u,o,W);
	}
}



void dfs2(int u,int fa,int o){
	for(int i=h[u];~i;i=e[i].nxt){
		int v=e[i].v;
		if(v==fa||v==w[u]) continue;
		dfs2(v,u,-1);             //消除轻儿子贡献
	}
	
	if(w[u]) dfs2(w[u],u,1);       //记录重儿子贡献
	
	calc(u,fa,1,w[u]);
	ans[u]=cc;
	
	if(o==-1){                 //记录重儿子贡献
		calc(u,fa,-1,0);
		mc=0;
		cc=0;
	}
	
}

int main(){
	init();
	int n;
	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,1);
	
	for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
	
	
	return 0;
}