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

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]);
}

标签：int,dfs,son,fa,gelral,clr,pson,CF600E,Lomsat
来源： https://www.cnblogs.com/codjjj/p/15934616.html