CF600E Lomsat gelral(樹上啟發式合併)
阿新 • • 發佈:2022-02-25
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]); }