CF600E Lomsat gelral(树上启发式合并)

Posted 2022-02-25 Codjjj

LINK

Code:

#include <iostream>
#include <cstring>
#include <algorithm>
#define int long long
typedef long long LL;

using namespace std;
const int N = 1e5+10,M=2*N;
int n,m;
int h[N],e[M],ne[M],idx;
void add(int a,int b)
    ne[idx]=h[a],e[idx]=b,h[a]=idx++;

int sz[N],son[N];
LL cnt[N],sum,mx;//记录当前节点为根子树的颜色数，最大颜色之和，最大颜色
LL ans[N];
int clr[N];
int dfs_son(int u,int fa)
    sz[u]=1;
    int pson=0;
    for(int i=h[u];~i;i=ne[i])
        int k=e[i];
        if(k==fa)continue;
        dfs_son(k,u);
        sz[u]+=sz[k];
        if(sz[k]>sz[pson])pson=k;
    
    son[u]=pson;
    return sz[u];

void calc(int u,int fa,int pson)//算所有轻儿子的贡献，不算u的重儿子子树
    cnt[clr[u]]++;
    if(cnt[clr[u]]>mx)mx=cnt[clr[u]],sum=clr[u];
    else if(cnt[clr[u]]==mx)sum+=clr[u];
    for(int i=h[u];~i;i=ne[i])
        int k=e[i];
        if(k==fa||k==pson)continue;
        calc(k,u,pson);
    

void remov(int u,int fa)//消除当前子树所有节点贡献
    int c=clr[u];
    cnt[c]--;
    for(int i=h[u];~i;i=ne[i])
        int k=e[i];
        if(k==fa)continue;
        remov(k,u);
    

void dfs(int u,int fa,int op)//op=1表示重儿子，0表示轻儿子
    for(int i=h[u];~i;i=ne[i])
        int k=e[i];
        if(k==fa||k==son[u])continue;//先跳过重儿子
        dfs(k,u,0);
    
    if(son[u])dfs(son[u],u,1);//最后遍历轻儿子
    calc(u,fa,son[u]);
    ans[u]=sum;
    if(!op)remov(u,fa),mx=sum=0;//记得清空

signed main()

    cin>>n;
    memset(h, -1, sizeof h);
    for(int i=1;i<=n;i++)cin>>clr[i];
    for(int i=1;i<n;i++)
        int a,b;
        cin>>a>>b;
        add(a,b),add(b,a);
    
    dfs_son(1,-1);
    dfs(1,-1,1);
    for(int i=1;i<=n;i++)printf("%lld ",ans[i]);