2018.08.28 codeforces600E（dsu on tree）

传送门

一道烂大街的dsu on tree板题。

感觉挺有趣的^_^

代码真心简单啊！

就是先处理轻儿子，然后处理重儿子，其中处理轻儿子后需要手动消除影响。

代码：

#include<bits/stdc++.h>
#define N 100005
#define ll long long
using namespace std;
inline int read(){
    int ans=0;
    char ch=getchar();
    while(!isdigit(ch))ch=getchar();
    while(isdigit(ch))ans=(ans<<3)+(ans<<1)+(ch^48),ch=getchar();
    return ans;
}
struct Node{int v,next;}e[N<<1];
int n,cnt,first[N],mp[N],tim[N],hson[N],siz[N],fa[N],maxv=0;
bool vis[N];
ll ans[N],sum=0;
inline void add(int u,int v){e[++cnt].v=v,e[cnt].next=first[u],first[u]=cnt;}
inline void dfs1(int p){
    siz[p]=1;
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(v==fa[p])continue;
        fa[v]=p,dfs1(v),siz[p]+=siz[v];
        if(siz[v]>siz[hson[p]])hson[p]=v;
    }
}
inline void modify(int p,int k){
    tim[mp[p]]+=k;
    if(k&&tim[mp[p]]>=maxv){
        if(tim[mp[p]]>maxv)sum=0,maxv=tim[mp[p]];
        sum+=mp[p];
    }
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(v==fa[p]||vis[v])continue;
        modify(v,k);
    }
}
inline void dfs2(int p,int used){
    cerr<<p<<'\n';
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(v!=hson[p]&&v!=fa[p])dfs2(v,0);
    }
    if(hson[p])dfs2(hson[p],1),vis[hson[p]]=1;
    modify(p,1),ans[p]=sum;
    if(hson[p])vis[hson[p]]=0;
    if(!used)modify(p,-1),maxv=sum=0;
}
int main(){
    n=read();
    for(int i=1;i<=n;++i)mp[i]=read();
    for(int i=1;i<n;++i){
        int u=read(),v=read();
        add(u,v),add(v,u);
    }
    dfs1(1),dfs2(1,0);
    for(int i=1;i<=n;++i)cout<<ans[i]<<' ';
    return 0;
}