CodeForces 600E. Lomsat gelral树上启发式合并

Posted 2021-03-02 bakacirno

传送门
好像大家都是拿这道题作为树上启发式合并的板子题。
树上启发式合并，英文是 dsu on tree，感觉还是中文的说法更准确，因为这个算法和并查集(dsu)没有任何关系。一般用来求解有根树的所有子树的统计问题。
根据轻重儿子的各种性质，可以证明这个算法的时间复杂度为 (O(nlogn))，虽然看起来暴力的不行，但是却是一个很高效的算法。
算法的核心其实就是对于每个节点，先计算轻儿子，再计算重儿子，把自己和轻儿子的所有贡献累计到重儿子上去，如果自己是轻儿子，就把贡献清空。
如果掌握了树链剖分和点分治，理解这个算法的流程还算挺简单的。

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <queue>
#define xx first
#define yy second
using namespace std;
typedef long long LL;
typedef pair<int,int> PII;
const int inf=0x3f3f3f3f;
const LL INF=0x3f3f3f3f3f3f3f3f;
const int N=3e5+10;
const int M=1e6+10;
int n,clr[N],siz[N],son[N],Son;
LL sum,maxc,cnt[N],ans[N];
vector<int> g[N];

void predfs(int u,int fa){
	siz[u]=1;
	for(int v:g[u]){
		if(v==fa) continue;
		predfs(v,u);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}

void calc(int u,int fa,int val){
	cnt[clr[u]]+=val;
	if(cnt[clr[u]]>maxc) maxc=cnt[clr[u]],sum=clr[u];
	else if(cnt[clr[u]]==maxc) sum+=clr[u];
	for(int v:g[u]){
		if(v==fa||v==Son) continue;
		calc(v,u,val);
	}
}

void dfs(int u,int fa,bool keep){
	//计算轻儿子
	for(int v:g[u]){
		if(v==fa||v==son[u]) continue;
		dfs(v,u,false);
	}
	//计算重儿子 
	if(son[u]) dfs(son[u],u,true),Son=son[u];
	//计算自己这颗子树的答案 
	calc(u,fa,1);
	ans[u]=sum;
	Son=0;
	//如果自己是轻儿子的话，就是清空自己对父亲的贡献 
	if(!keep) calc(u,fa,-1),sum=maxc=0;
}

int main(){
	scanf("%d",&n);
	for(int i=1;i<=n;i++) scanf("%d",&clr[i]);
	for(int i=1,u,v;i<n;i++){
		scanf("%d%d",&u,&v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	predfs(1,1);
	dfs(1,1,1);
	for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
	return 0;
}