CF600E Lomsat gelral
-
\(dsu\ on\ tree\) 入门例题.讲解.
- 注意 \(dfs\) 中各个操作的顺序,这样能保证清空轻儿子贡献时不会对重儿子造成影响.
#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define mp make_pair
#define pii pair<int,int>
inline int read()
{
int x=0;
bool pos=1;
char ch=getchar();
for(;!isdigit(ch);ch=getchar())
if(ch=='-')
pos=0;
for(;isdigit(ch);ch=getchar())
x=x*10+ch-'0';
return pos?x:-x;
}
const int MAXN=1e5+10;
vector<int> sons[MAXN];
int n;
int mxson[MAXN],siz[MAXN],col[MAXN],cnt[MAXN];
ll sum,ans[MAXN];
void getsiz(int u,int fa)
{
siz[u]=1;
int tot=sons[u].size();
for(int i=0;i<tot;++i)
{
int v=sons[u][i];
if(v!=fa)
{
getsiz(v,u);
if(siz[v]>siz[mxson[u]])
mxson[u]=v;
siz[u]+=siz[v];
}
}
}
int mxcnt,bigson;
void add(int u,int fa,int val)
{
cnt[col[u]]+=val;
if(cnt[col[u]]>mxcnt)
mxcnt=cnt[col[u]],sum=col[u];
else if(cnt[col[u]]==mxcnt)
sum+=col[u];
int tot=sons[u].size();
for(int i=0;i<tot;++i)
{
int v=sons[u][i];
if(v!=fa && v!=bigson)
add(v,u,val);
}
}
void dfs(int u,int fa,bool keep)
{
int tot=sons[u].size();
for(int i=0;i<tot;++i)
{
int v=sons[u][i];
if(v!=fa && v!=mxson[u])
dfs(v,u,false);
}
if(mxson[u])
dfs(mxson[u],u,true),bigson=mxson[u];
add(u,fa,1);
ans[u]=sum;
bigson=0;
if(keep==false)
add(u,fa,-1),sum=0,mxcnt=0;
}
int main()
{
n=read();
for(int i=1;i<=n;++i)
col[i]=read();
for(int i=1;i<n;++i)
{
int u=read(),v=read();
sons[u].push_back(v);
sons[v].push_back(u);
}
getsiz(1,0);
dfs(1,0,true);
for(int i=1;i<=n;++i)
printf("%I64d ",ans[i]);
return 0;
}