codeforces600E Lomsat gelral -- 树上启发式合并

时间:2021-08-06 09:59:09

显然加入和删除点可以 O(1) 维护。先做一次树剖,然后从根开始 dfs 一遍。 dfs 到一个点时加入答案。对于每个点如果是轻儿子退出时删除,否则不删除。统计答案时暴力 dfs 一遍所有轻儿子。
因为一个点到根路径上只有 O(logn) 条轻边,所以最多只会被统计 O(logn) 次。时间复杂度为 O(nlogn)

代码:

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
#define N 100010
#define ll long long
vector<int>g[N];
ll Sum[N],Ans[N];
int i,j,k,n,m,M,sz[N],s[N],a[N],Res,c[N],x,y;
inline int Max(int x,int y){
return x<y?y:x;
}
inline void Dfs1(int x,int y){
sz[x]=1;
for(int i=0;i<g[x].size();i++)
if(g[x][i]!=y){
Dfs1(g[x][i],x);
sz[x]+=sz[g[x][i]];
if(sz[g[x][i]]>sz[s[x]])s[x]=g[x][i];
}
}
inline void Add(int x){
Sum[++c[x]]+=x;
Sum[c[x]-1]-=x;
Res=Max(Res,c[x]);
}
inline void Del(int x){
Sum[c[x]]-=x;
if(!Sum[c[x]]&&Res==c[x])Res--;
Sum[--c[x]]+=x;
}
inline void Dfs2(int x,int y,bool d){
for(int i=0;i<g[x].size();i++)
if(g[x][i]!=y)Dfs2(g[x][i],x,d);
if(d)Add(a[x]);else Del(a[x]);
}
inline void Dfs(int x,int y,bool d){
for(int i=0;i<g[x].size();i++)
if(g[x][i]!=y&&g[x][i]!=s[x])Dfs(g[x][i],x,0);
if(s[x])Dfs(s[x],x,1);Add(a[x]);
for(int i=0;i<g[x].size();i++)
if(g[x][i]!=y&&g[x][i]!=s[x])Dfs2(g[x][i],x,1);
Ans[x]=Sum[Res];
if(!d){
for(int i=0;i<g[x].size();i++)
if(g[x][i]!=y)Dfs2(g[x][i],x,0);
Del(a[x]);
}
}
int main(){
scanf("%d",&n);
for(i=1;i<=n;i++)scanf("%d",&a[i]);
for(i=1;i<n;i++)scanf("%d%d",&x,&y),g[x].push_back(y),g[y].push_back(x);
Dfs1(1,0);
Dfs(1,0,1);
for(i=1;i<=n;i++)printf("%I64d ",Ans[i]);
return 0;
}