树上启发式合并入门-Lomsat gelral

简述：

树上启发式合并又叫 dsu on tree ，是解决可以离线的子树问题的一种方法。
优点：时间复杂度优秀$O(nlogn)$，在一些可以使用 dfs 序 + 莫队 算法$O(n\sqrt{n})$的题目中有很强的优越性，相比线段树码量要小的多。
核心思想：树上启发式合并利用重链剖分的性质优化子树贡献的计算。
题目链接

思路：

其实是树上启发式合并的入门题目，考虑小的集合向大的集合合并，大的集合就是重儿子。

1. 找重儿子这里和树链剖分是一样的
2. 对每个以当前为根的子树，都先走他的轻儿子，再走他的重儿子（和重链剖分正好相反，树上启发式合并重点之一）。
3. 在走每个轻儿子时都要计算并记录答案，需要开一个数组记录颜色个数，由于 dfs 的特性在访问另外的儿子时需要清空数组（简单解释下，这个问题在 dfs 解决子树问题时很常见，数组是一维的，父的信息是由所有子得来的，在计算子的答案产生的信息会影响其他子的计算）
4. 最后计算重儿子并记录答案，不清空数组（树上启发式合并重点之一，此时子的答案算完了），暴力加入所有轻儿子的信息，计算父节点。
5. 输出答案 for 一遍。

乱搞下时间复杂度证明：

每次是小的集合向大的集合合并，时间长在了轻儿子的暴力，一个节点可能进出多次，次数是它作为轻边的次数，而由于重链剖分的性质，任意一个点往上跳时所经过的重链数量不超过 $O(logn)$，所以重链交界处（轻边）的数量也不会超过 $O(logn)$ ,因此每个点的统计次数也为 $O(logn)$。总时间复杂度$O(nlogn)$。
关于重链剖分的性质证明建议百度。。

Code:

莫队+dfs 序也可以写，但懒得写了。

#include <iostream>
#include <map>
#include <set>
#include <queue>
#include <stack>
#include <algorithm>
#include <vector>
#include <string>
#include <iomanip>
#include <cmath>
#include <ctime>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <climits>
//#include <unordered_map>
#define guo312 std::ios::sync_with_stdio(false), cin.tie(0), cout.tie(0)
#define ll long long
#define Inf LONG_LONG_MAX
#define inf INT_MAX
#define endl "\n"
#define PI 3.1415926535898
using namespace std;
const int N=1e5+10;
int n,col[N];
ll ans[N],cnt[N];
vector<int> ve[N];

int son[N],size[N];
void dfs1(int now,int last){  // 找重儿子
	size[now]=1;
	int maxn=0;
	for(auto it:ve[now]){
		if(it!=last){
			dfs1(it,now);
			size[now]+=size[it];
			if(size[it]>maxn){
				maxn=size[it],son[now]=it;
			}
		}
	}
}

ll maxn=0,sum=0;
queue<int> p;  // 记录那些颜色近队方便清空
int flag[N];

void insert(int x){  // 插入一个颜色
	cnt[col[x]]++;
	if(cnt[col[x]]==maxn) sum+=col[x];
	else if(cnt[col[x]]>maxn) maxn=cnt[col[x]],sum=col[x];
	if(flag[col[x]]==0) p.push(col[x]),flag[col[x]]=1;
}

void add(int now,int last){  // 加入以now为根的所有子树
	insert(now);
	for(auto it:ve[now]){
		if(it!=last){
			add(it,now);
		}
	}
}

void clear(){  // 清空信息
	while(!p.empty()){
		int s=p.front(); p.pop();
		cnt[s]=0,flag[s]=0;
	}
	sum=0,maxn=0;
}

void dfs2(int now,int last){  // 树上启发式合并主体
	for(auto it:ve[now]){
		if(it!=last&&it!=son[now]){
			dfs2(it,now);  // 走轻儿子
			clear();  // 清空
		}
	}
	if(son[now]) dfs2(son[now],now);  // 存在重儿子就走（叶子没有重儿子）
	for(auto it:ve[now]){  // 暴力加轻儿子
		if(it!=last&&it!=son[now]){
			add(it,now);
		}
	}
	insert(now);  // 父节点不要忘
	ans[now]=sum;
}

int main(){
guo312;
	cin>>n;
	for(int i=1;i<=n;i++){
		cin>>col[i];
	}
	for(int i=1;i<n;i++){
		int s1,s2; cin>>s1>>s2;
		ve[s1].push_back(s2);
		ve[s2].push_back(s1);
	}
	dfs1(1,0);
	dfs2(1,0);
	for(int i=1;i<=n;i++){
		cout<<ans[i]<<" ";
	}
	return 0;
}