DSU on tree 通常用来解决不带修树上子树问题。
主要思想:
- 剖分。
- 先搜轻儿子,记录轻儿子子树的答案,删去轻儿子的贡献。
- 搜重儿子,记录重儿子子树的答案,保留重儿子的贡献。
- 回溯,重新搜轻儿子,把轻儿子子树的贡献加上,构成本子树的答案。
#include<bits/stdc++.h>
using namespace std;
#define ll long long
const int N=1e5+10;
struct edge
{
int u,v,nxt;
}e[N<<1];
int head[N],tot;
void add(int u,int v)
{
e[++tot]=edge{u,v,head[u]};
head[u]=tot;
}
int n;
int col[N];
int siz[N],son[N];
void dfs1(int u,int f)
{
siz[u]=1;
for(int i=head[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==f) continue;
dfs1(v,u);
siz[u]+=siz[v];
if(siz[v]>siz[son[u]]) son[u]=v;
}
}
int cnt[N];
ll ans[N];
int maxn;
ll sum;
void add(int u,int f,int s)
{
cnt[col[u]]++;
if(cnt[col[u]]>maxn) maxn=cnt[col[u]],sum=col[u];
else if(cnt[col[u]]==maxn) sum+=col[u];
for(int i=head[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==f||v==s) continue;
add(v,u,s);
}
}
void sub(int u,int f)
{
cnt[col[u]]--;
for(int i=head[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==f) continue;
sub(v,u);
}
}
void dfs2(int u,int f,int s)
{
for(int i=head[u];i;i=e[i].nxt)
{
int v=e[i].v;
if(v==f||v==son[u]) continue;
dfs2(v,u,0);
}
if(son[u]) dfs2(son[u],u,1);
add(u,f,son[u]);
ans[u]=sum;
if(!s) sub(u,f),sum=maxn=0;
}
int main()
{
scanf("%d",&n);
for(int i=1;i<=n;i++) scanf("%d",&col[i]);
for(int i=1;i<n;i++)
{
int u,v;
scanf("%d%d",&u,&v);
add(u,v),add(v,u);
}
dfs1(1,0),dfs2(1,0,0);
for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
return 0;
}