CodeForces 600E. Lomsat gelral【树上启发式合并】
传送门
好像大家都是拿这道题作为树上启发式合并的板子题。
树上启发式合并,英文是 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;
}