CodeForces Lomsat gelral（DSU on tree 树上启发式合并）

传送门：https://codeforces.com/problemset/problem/600/E

题意：

给一颗n个点的有根树，每个点有一种颜色，一颗子树中出现次数最多的颜色称为主导颜色，主导颜色可能有多个，求这棵树的每个结点的对应子树的主导颜色之和

思路：

DSU on tree 树上启发式合并，先记录一下板子
··········································································································
先树链剖分找重儿子，重儿子就是最大子树的根节点
dfs每个节点，暴力计算所有轻儿子的贡献，并合并到重儿子上去
注意，轻儿子的贡献在计算之后必须要删除，不然可能会MLE，
如果都不保留，可能会TLE，所以需要平衡一下时间和空间，我们保留重儿子的贡献即可，这也是DSU on tree的意义

时间复杂度O(nlogn)，以下分析来自Yveh的[trick]dsu on tree
在这种做法中，我们先进行树链剖分。
dfs的时候，首先dfs节点x的轻儿子，暴力消去影响，再dfs节点x的下一个轻儿子，依此类推。
然后dfs节点x的重儿子，无需消去影响。
在最后，我们为了统计x，再将x轻儿子的贡献加回来。
看起来很暴力，但是实际上它的时间复杂度是O(nlogn)的，跑得飞快。
可以这么考虑：
只有dfs到轻边时，才会将轻边的子树中合并到上一级的重链，树链剖分将一棵树分割成了不超过logn条重链。
每一个节点最多向上合并logn次，单次修改复杂度O(1)。
所以整体复杂度是O(nlogn)的。

AC代码：

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=1e5+10;
struct Edge {
	int to,next;
} e[maxn<<1];
int n,u,v,tot=1,head[maxn],siz[maxn],son[maxn];
int col[maxn],cnt[maxn],flag,maxc;
ll ans[maxn],sum;
inline void addEdge(int u,int v) {
	e[tot].to = v;
	e[tot].next = head[u];
	head[u]=tot++;
}
//重链剖分找重儿子
void dfs(int u,int f) {
	siz[u]=1;
	for(int i=head[u]; i; i=e[i].next) {
		int v=e[i].to;
		if(v==f) continue;
		dfs(v,u);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}
//统计某结点及其所有轻儿子的贡献
void count(int u,int f,int val) {
	cnt[col[u]] += val;//val为正负用来控制是加贡献还是删除贡献
	if(cnt[col[u]]>maxc) {
		maxc=cnt[col[u]];
		sum=col[u];
	} else if(cnt[col[u]]==maxc) { //两个颜色数量相同都是主导颜色都得算
		sum+=col[u];
	}
	//排除被标记的重儿子，统计其他儿子子树信息
	for(int i=head[u]; i; i=e[i].next) {
		int v=e[i].to;
		if(v==f || v==flag) continue;
		count(v,u,val);
	}
}
//DSU on tree 的板子
void dfs(int u,int f,bool keep) {
	//第一步：搞轻儿子及其子树 算其答案 删贡献
	for(int i=head[u]; i; i=e[i].next) {
		int v=e[i].to;
		if(v==f || v==son[u]) continue;
		dfs(v,u,false);
	}

	//第二步：搞重儿子及其子树算其答案 不删贡献
	if(son[u]) {
		dfs(son[u],u,true);
		flag=son[u];
	}

	//第三步：暴力统计u及其所有轻儿子的贡献合并到刚算出的重儿子信息里
	count(u,f,1);
	flag = 0;
	ans[u] = sum;

	//把需要删除贡献的删一删
	if(!keep) {
		count(u,f,-1);
		sum=maxc=0;
	}

}
int main() {
	cin>>n;
	for(int i=1; i<=n; i++) cin>>col[i];
	for(int i=1; i<n; i++) {
		cin>>u>>v;
		addEdge(u,v),addEdge(v,u);
	}
	dfs(1,0);
	dfs(1,0,0);
	for(int i=1; i<=n; i++) {
		cout<<ans[i]<<" ";
	}
}