CF Lomsat gelral（启发式合并 dsu on tree）

洛谷链接：Lomsat gelral - 洛谷

题意：

• 有一棵 nn 个结点的以 1 号结点为根的有根树。
• 每个结点都有一个颜色，颜色是以编号表示的， i 号结点的颜色编号为 c[i]​。
• 如果一种颜色在以 x 为根的子树内出现次数最多，称其在以 x 为根的子树中占主导地位。显然，同一子树中可能有多种颜色占主导地位。
• 你的任务是对于每一个 i∈[1,n]，求出以 i 为根的子树中，占主导地位的颜色的编号和。

这里要用到启发式合并，我对启发式合并的理解如下：

本题要求“所有子树的信息”。 我们由树形dp的思想知道，一个结点的子树的信息是其所有儿子的子树信息+自己本身的信息。但是这里我们遇到一个问题，就是无法O(1)地把信息合并，假设每合并两个信息是O(n)的，它有m个儿子，那么完成该结点的信息合并就已经需要O(nm)了，这就不行。

那么我们就考虑这样：让该结点继承它一个儿子的子树信息，然后暴力统计其他信息（包括自己的信息和其他儿子的子树信息），发现这样更快。那么我们选择哪个儿子进行继承呢？重儿子，因为重儿子暴力统计起来最慢。所以启发式合并需要用到轻重链剖分。

按照这个思路，启发式合并的一般步骤如下：

1. 跑个dfs1来记录子树大小，高度，重儿子等信息。

2.再跑dfs2进行答案计算，对于每个结点，继承其重儿子的信息，然后暴力统计其他信息。

code：

#include<bits/stdc++.h>
using namespace std;
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define int long long
const int N = 1e5+5;
int n, a[N], sz[N], son[N];
int cnt[N], ans[N], Mx, Son, Ans;
vector<int> g[N];
 
void dfs1(int u,int fa){
    sz[u] = 1;
    for(int v:g[u]){
        if(v==fa) continue;
        dfs1(v,u);
        sz[u] += sz[v];
        if(sz[v] > sz[son[u]]) son[u] = v; //轻重链剖分，记录重儿子
    }
}
void add(int u,int fa,int val){
    cnt[a[u]] += val; //统计当前的贡献
    if(cnt[a[u]] > Mx) Mx=cnt[a[u]], Ans=a[u]; //这两行是更新最优解信息
    else if(cnt[a[u]]==Mx) Ans+=a[u];
    for(int v:g[u]){ //递归处理子树
        if(v==fa || v==Son) continue;
        add(v,u,val);
    }
}
void dfs2(int u,int fa,int op){ //op=0表示不用保留信息，op=1表示要保留
    for(int v:g[u]){
        if(v==fa) continue;
        if(v!=son[u]) dfs2(v,u,0); //先dfs统计轻儿子的信息，然后扔掉信息
    }
    if(son[u]) dfs2(son[u],u,1), Son=son[u]; //再dfs统计重儿子的信息并保留信息。并且记录当前重儿子，用于后面add
 
    add(u,fa,1); //继承重儿子信息后，暴力统计剩余部分
    Son = 0; //重置当前重儿子
    ans[u] = Ans; //记录当前答案
    if(!op) add(u,fa,-1), Ans=0, Mx=0; //删掉信息
}
inline void solve(){
    cin>>n; FOR(i,1,n) cin>>a[i];
    FOR(i,1,n-1){
        int x,y; cin>>x>>y;
        g[x].push_back(y); g[y].push_back(x);
    }
    dfs1(1,0); dfs2(1,0,0);
    FOR(i,1,n) cout<<ans[i]<<' ';
}
signed main(){
    ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
    int T=1; while(T--) solve();
}

---

再加一道dsu on tree的例题：Problem - 1009F - Codeforces

code：

#include <bits/stdc++.h>
using namespace std;
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define int long long
const int N = 1e6+5;
int n,sz[N],son[N],d[N];
int Son, Mx, Ans, cnt[N]; //cnt[i]表示当前子树中深度为i的结点数量，Mx记录max(cnt[i])，Ans记录对应的Mx对应的i，Son是当前重儿子，add中不用算它的信息
int ans[N];
vector<int> g[N];

void add(int u,int fa,int val){
	cnt[d[u]] += val; //u结点自身贡献
	if(cnt[d[u]]>Mx || (cnt[d[u]]==Mx && d[u]<Ans)) Mx=cnt[d[u]], Ans=d[u]; //看看当前答案要不要改变
	for(int v:g[u]){
		if(v==fa || v==Son) continue; //跳过重儿子和父亲
		add(v,u,val);
	}
}
void dfs1(int u,int fa){
	sz[u]=1, d[u]=d[fa]+1;
	for(int v:g[u]){
		if(v==fa) continue;
		dfs1(v,u);
		sz[u] += sz[v];
		if(sz[v] > sz[son[u]]) son[u]=v;
	}
}
void dfs2(int u,int fa,int op){ //op=0表示不保留信息，op=1表示保留
	for(int v:g[u]){
		if(v==fa || v==son[u]) continue;
		dfs2(v,u,0); //暴力算轻儿子信息，并且扔掉信息
	}
	if(son[u]) dfs2(son[u],u,1), Son=son[u]; //继承重儿子信息，并且记录重儿子编号
	add(u,fa,1); //暴力统计剩下的信息
	Son = 0; //重置重儿子
	ans[u] = Ans-d[u]; //记录当前答案
	if(!op) add(u,fa,-1), Ans=Mx=0; //如果不是重儿子，删掉对应信息
}
void solve(){
	cin>>n;
	FOR(i,1,n-1){
		int u,v; cin>>u>>v;
		g[u].push_back(v), g[v].push_back(u);
	}
	dfs1(1,0);
	dfs2(1,0,1);

	FOR(i,1,n) cout<<ans[i]<<'\n';
}
signed main(){
	ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
    int T=1; //cin>>T;
	while(T--) solve();
}