600E.Lomsat gelral (树上启发式合并)
给出一个n个点的树
求出树中每个子树出现次数最多的所有颜色的编号的和。
题解:
搜索的过程中,重儿子所在的子树搜一遍,记录颜色的出现情况
然后把轻儿子往重儿子上合并
#include<bits/stdc++.h>
using namespace std;
const int maxn=1e5+100;
int c[maxn];
vector<int> g[maxn];
long long ans[maxn];
int f[maxn];//当前节点的子树内每种颜色的出现次数
int L[maxn],R[maxn],id[maxn];
int dep[maxn];
int size[maxn];
int tot;
int son[maxn];
int Max[maxn];
int n;
void dfs1 (int x,int pre) {
dep[x]=dep[pre]+1;
size[x]=1;
L[x]=++tot;
id[tot]=x;
for (int y:g[x]) {
if (y==pre) continue;
dfs1(y,x);
size[x]+=size[y];
if (size[son[x]]<size[y]) son[x]=y;
}
R[x]=tot;
}
void cal (int x,int pre) {
//计算x的答案
//在cal函数里
//重儿子的f状态已经保存
f[c[x]]++;
if (f[c[x]]>Max[x]) {
Max[x]=f[c[x]];
ans[x]=c[x];
}
else if (f[c[x]]==Max[x]) {
ans[x]+=c[x];
}
for (int y:g[x]) {
//遍历除重儿子外的子树
//把答案往重儿子上合并
if (y==son[x]) continue;
if (y==pre) continue;
for (int j=L[y];j<=R[y];j++) {
int z=id[j];
f[c[z]]++;
if (f[c[z]]>Max[x]) {
Max[x]=f[c[z]];
ans[x]=c[z];
}
else if (f[c[z]]==Max[x]) {
ans[x]+=c[z];
}
}
}
// Max[x]=0;
// ans[x]=0;
// for (int i=1;i<=n;i++) Max[x]=max(Max[x],f[i]);
// for (int i=1;i<=n;i++) if (f[i]==Max[x]) ans[x]+=i;
}
void dfs2 (int x,int pre,int kp) {
for (int y:g[x]) {
if (y==son[x]) continue;
if (y==pre) continue;
dfs2(y,x,0);
}
if (son[x]) {
dfs2(son[x],x,1);
ans[x]=ans[son[x]];
Max[x]=Max[son[x]];
}
cal(x,pre);
if (!kp) for (int i=L[x];i<=R[x];i++) f[c[id[i]]]=0;
}
int main () {
scanf("%d",&n);
for (int i=1;i<=n;i++) scanf("%d",c+i);
for (int i=1;i<n;i++) {
int x,y;
scanf("%d%d",&x,&y);
g[x].push_back(y);
g[y].push_back(x);
}
dfs1(1,0);
dfs2(1,0,1);
for (int i=1;i<=n;i++) printf("%lld ",ans[i]);
}