树上启发式合并
题目:一棵树有n个结点,每个结点都是一种颜色,每个颜色有一个编号,求树中每个子树的最多的颜色编号的和。
树上启发式合并真神奇,时间复杂度只有O(nlogn)
看下思路
第一步:我先像树剖那样,跑出重儿子。
第二步:我们dfs这个树,比如说,我们现在跑到了u节点。首先优先跑所有轻儿子,采用尾递归的方式,统计轻儿子的答案。期间维护一个通数组cnt,存的是以当前轻儿子为根的子树的颜色数量。跑完后将答案赋予这个轻儿子节点,然后清空cnt数组,开始跑下一个轻儿子。
第三步:轻儿子跑完后,我们开始跑重儿子。跑完后,注意:我们不在清空cnt数组,然后逐个再跑一遍各个轻儿子。将结果累计到cnt里面,这里就体现了我们所受的启发:小往大合并。
第四步:经第三步,我们cnt数组就存入了所有子树信息,然后结合u节点自身的信息,将这个答案赋予u就行啦。
#include<cstdio>
#define int long long
const int maxn=1e5+5;
int ver[maxn<<1],head[maxn],next[maxn<<1];
int tot=1;
void add(int x,int y){
ver[++tot]=y;
next[tot]=head[x];
head[x]=tot;
}
int siz[maxn],son[maxn],col[maxn],cnt[maxn],ans[maxn];
void dfs_bson(int u,int f){
siz[u]=1;
for(int i=head[u];i;i=next[i]){
int y=ver[i];
if(y==f) continue;
dfs_bson(y,u);
siz[u]+=siz[y];
if(siz[y]>siz[son[u]])
son[u]=y;
}
}
int sum,flag,maxx;
void _count(int u,int f,int val){
cnt[col[u]]+=val;
if(cnt[col[u]]>maxx){
maxx=cnt[col[u]];
sum=col[u];
}
else if(cnt[col[u]]==maxx) sum+=col[u];
for(int i=head[u];i;i=next[i]){
int y=ver[i];
if(y==f||y==flag) continue;
_count(y,u,val);
}
}
void dfs(int u,int f,bool _flag){
for(int i=head[u];i;i=next[i]){
int y=ver[i];
if(y==f||y==son[u]) continue;
dfs(y,u,0);
}
if(son[u]){
dfs(son[u],u,1);
flag=son[u];
}
_count(u,f,1);
flag=0;
ans[u]=sum;
if(!_flag){
_count(u,f,-1);
sum=maxx=0;
}
}
signed main(){
int n;
scanf("%lld",&n);
for(int i=1;i<=n;i++)
scanf("%lld",&col[i]);
for(int i=1;i<n;i++){
int x,y;
scanf("%lld%lld",&x,&y);
add(x,y);
add(y,x);
}
dfs_bson(1,0);
dfs(1,0,0);
for(int i=1;i<=n;i++){
printf("%lld",ans[i]);
if(i!=n) printf(" ");
}
return 0;
}