Codeforces 600E. Lomsat gelral [dsu on tree]

题解传送门

新姿势

dsu on tree 模板题

#include<cstdio>
#include<vector>
#include<algorithm>
#define N 100005
#define LL long long
using namespace std;

vector<int>e[N];
int n,c[N],sz[N],son[N],num[N],mx;
LL ans[N],sum;
bool vis[N];

inline int read(){
    int a=0;char f=1,c=getchar();
    while(c<'0'||c>'9'){if(c=='-')f=-1;c=getchar();}
    while(c>='0'&&c<='9'){a=a*10+c-'0';c=getchar();}
    return a*f;
}

inline void addedge(int x,int y){e[x].push_back(y),e[y].push_back(x);}

void dfs(int x,int fa){
    sz[x]=1;
    for(int i=0;i<e[x].size();++i)
        if(e[x][i]!=fa){
            dfs(e[x][i],x);
            sz[x]+=sz[e[x][i]];
            if(sz[e[x][i]]>sz[son[x]]) son[x]=e[x][i];
        }
}

inline void change(int x,int fa,int v){
    num[c[x]]+=v;
    if(num[c[x]]==mx) sum+=1ll*c[x];
    if(num[c[x]]>mx) mx=num[c[x]],sum=c[x];
    for(int i=0;i<e[x].size();++i)
        if(e[x][i]!=fa&&!vis[e[x][i]]) change(e[x][i],x,v);
}

void dfs(int x,int fa,int opt){
    for(int i=0;i<e[x].size();++i)
        if(e[x][i]!=fa&&e[x][i]!=son[x]) dfs(e[x][i],x,0);
        if(son[x]) dfs(son[x],x,1),vis[son[x]]=1;
    change(x,fa,1);
    ans[x]=sum;
    if(son[x]) vis[son[x]]=0;
    if(!opt) change(x,fa,-1),sum=mx=0;      
}

int main(){
    n=read();
    for(int i=1;i<=n;++i) c[i]=read();
    for(int i=1;i<n;++i){
        int x=read(),y=read();
        addedge(x,y);
    }
    dfs(1,0);
    dfs(1,0,1);
    for(int i=1;i<=n;++i) printf("%lld%s",ans[i],i==n?"\n":" ");
    return 0;
}