CF 600E 树上众数 启发式合并

顺便实现了在每次增减1的情况下,数组最大值的维护。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
struct edge{
    int to,next;
}e[200003];
int head[100003];
int cnt;
void init(){
    memset(head,-1,sizeof(head));
    cnt=0;
}
void add(int u,int v){
    e[cnt].to=v;
    e[cnt].next=head[u];
    head[u]=cnt++;
}
int mc[100003];
int c[100003];
int n;
ll ans[100003];
int siz[100003];
int tim[100003];
ll grp[100003];
int mx;
ll sum;
void dfs(int u,int fa){
    siz[u]=1;
    mc[u]=-1;
    for(int i=head[u];~i;i=e[i].next){
        int v=e[i].to;
        if(v==fa)continue;
        dfs(v,u);
        siz[u]+=siz[v];
        if(mc[u]==-1||siz[mc[u]]<siz[v])mc[u]=v;
    }
}
void add(int u){
    int p=++tim[c[u]];
    if(p==1){
        grp[1]+=c[u];
    }
    else {
        grp[p]+=c[u];
        grp[p-1]-=c[u];
    }
    if(p>=mx){
        mx=p;
    }
}
void del(int u){
    int p=--tim[c[u]];
    if(p==0){
        grp[1]-=c[u];
    }
    else {
        grp[p]+=c[u];
        grp[p+1]-=c[u];
    }
    if(!grp[mx]){
        mx--;
    }
}
void addtree(int u,int fa){
    add(u);
    for(int i=head[u];~i;i=e[i].next){
        int v=e[i].to;
        if(v==fa)continue;
        addtree(v,u);
    }
}
void deltree(int u,int fa){
    del(u);
    for(int i=head[u];~i;i=e[i].next){
        int v=e[i].to;
        if(v==fa)continue;
        deltree(v,u);
    }
}
void dfs2(int u,int fa){
    if(mc[u]==-1){
        add(u);
        ans[u]=grp[mx];
        return;
    }
    for(int i=head[u];~i;i=e[i].next){
        int v=e[i].to;
        if(v==fa||v==mc[u])continue;
        dfs2(v,u);
        deltree(v,u);
    }
    dfs2(mc[u],u);
    for(int i=head[u];~i;i=e[i].next){
        int v=e[i].to;
        if(v==fa||v==mc[u])continue;
        addtree(v,u);
    }
    add(u);
    ans[u]=grp[mx];
    return ;
}
int main(){
    scanf("%d",&n);
    for(int i=1;i<=n;i++){
        scanf("%d",&c[i]);
    }
    init();
    for(int i=1;i<n;i++){
        int u,v;
        scanf("%d%d",&u,&v);
        add(u,v);
        add(v,u);
    }
    dfs(1,0);
    mx=0;
    sum=0;
    dfs2(1,0);
    for(int i=1;i<=n;i++)printf("%lld%c",ans[i],i==n?'\n':' ');
}

猜你喜欢

转载自blog.csdn.net/MrBird_to_fly/article/details/80876133