1. 程式人生 > 其它 >CF600E Lomsat gelral(樹上啟發式合併)

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