CF 600E. Lomsat gelral(dsu on tree)

Posted 2021-01-30 sdfzsyq

解题思路

　　(dsu) (on) (tree)的模板题。暴力而优雅的算法，轻儿子的信息暴力清空，重儿子的信息保留，时间复杂度(O(nlogn))

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<set>

using namespace std;
const int MAXN = 100005;
typedef long long LL;

inline int rd(){
    int x=0,f=1;char ch=getchar();
    while(!isdigit(ch)) f=ch==‘-‘?0:1,ch=getchar();
    while(isdigit(ch)) x=(x<<1)+(x<<3)+ch-‘0‘,ch=getchar();
    return f?x:-x;
}

int n,c[MAXN],head[MAXN],cnt,to[MAXN<<1],nxt[MAXN<<1],num[MAXN],sum[MAXN];
int siz[MAXN],son[MAXN],fa[MAXN],Num;
LL ans[MAXN],Max[MAXN],Ans;

inline void add(int bg,int ed){
    to[++cnt]=ed,nxt[cnt]=head[bg],head[bg]=cnt;
}

inline void update(int x){
    sum[num[c[x]]]--;Max[num[c[x]]]-=c[x];
    num[c[x]]++;sum[num[c[x]]]++;Max[num[c[x]]]+=c[x];
    if(Num<=num[c[x]]) Num=num[c[x]];Ans=Max[Num];
}

inline void Clear(int x){
    sum[num[c[x]]]--;Max[num[c[x]]]-=c[x];
    num[c[x]]--;sum[num[c[x]]]++;Max[num[c[x]]]+=c[x];
    if(Num==num[c[x]]+1 && !sum[num[c[x]]+1]) Num=num[c[x]];
    Ans=Max[Num];
}

void bfs(int x){
    update(x);
    for(int i=head[x];i;i=nxt[i])
        if(to[i]!=fa[x]) bfs(to[i]);
}

void clear(int x){
    Clear(x);
    for(int i=head[x];i;i=nxt[i])
        if(to[i]!=fa[x]) clear(to[i]);
}

void dfs1(int x,int f){
    fa[x]=f;siz[x]=1;int maxson=-1,u;
    for(int i=head[x];i;i=nxt[i]){
        u=to[i];if(u==f) continue;
        dfs1(u,x);siz[x]+=siz[u];
        if(siz[u]>maxson) maxson=siz[u],son[x]=u;
    }
}

void dfs2(int x){
    for(int i=head[x];i;i=nxt[i]){
        int u=to[i];if(u==fa[x] || u==son[x]) continue;
        dfs2(u);clear(u);
    }
    if(son[x]) dfs2(son[x]);
    for(int i=head[x];i;i=nxt[i]){
        int u=to[i];if(u==fa[x] || u==son[x]) continue;
        bfs(u);
    }update(x);
    ans[x]=Ans;
}

int main(){
    n=rd();int x,y;
    for(int i=1;i<=n;i++) c[i]=rd();
    for(int i=1;i<n;i++){
        x=rd(),y=rd();
        add(x,y);add(y,x);
    }dfs1(1,0);dfs2(1);
    for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
    return 0;
}