题目链接:点此跳转
题目大意:一棵树有n个结点,每个结点都是一种颜色,每个颜色有一个编号,求树中每个子树的最多的颜色编号的和。
解题思路:这道是树启的模板题,我们先跑一遍dfs处理出轻重儿子,然后每次保留重儿子的贡献,消除轻儿子的贡献,这样时间复杂度可以降到log级别。
Code:
#include<iostream>
#include<cstring>
#include<queue>
using namespace std;
typedef long long ll;
const int maxn=1e5+7;
struct edge{
int v,nxt;
}e[maxn<<1];
int h[maxn],cntt;
int col[maxn],cnt[maxn],sz[maxn],w[maxn],n;
ll ans[maxn],mc,cc;
void init(){
memset(h,-1,sizeof h);
cntt=0;
}
void add(int u,int v){
e[cntt].v=v;
e[cntt].nxt=h[u];
h[u]=cntt++;
}
void dfs1(int u,int fa){ //处理轻重儿子
sz[u]=1;
for(int i=h[u];~i;i=e[i].nxt){
int v=e[i].v;
if(v==fa) continue;
dfs1(v,u);
if(sz[v]>sz[w[u]]) w[u]=v;
sz[u]+=sz[v];
}
}
void calc(int u,int fa,int o,int W){ //计算结果
cnt[col[u]]+=o;
if(cnt[col[u]]==mc) cc+=col[u];
else if(cnt[col[u]]>mc) mc=cnt[col[u]],cc=col[u];
for(int i=h[u];~i;i=e[i].nxt){
int v=e[i].v;
if(v==fa||v==W) continue;
calc(v,u,o,W);
}
}
void dfs2(int u,int fa,int o){
for(int i=h[u];~i;i=e[i].nxt){
int v=e[i].v;
if(v==fa||v==w[u]) continue;
dfs2(v,u,-1); //消除轻儿子贡献
}
if(w[u]) dfs2(w[u],u,1); //记录重儿子贡献
calc(u,fa,1,w[u]);
ans[u]=cc;
if(o==-1){ //记录重儿子贡献
calc(u,fa,-1,0);
mc=0;
cc=0;
}
}
int main(){
init();
int n;
scanf("%d",&n);
for(int i=1;i<=n;i++) scanf("%d",&col[i]);
for(int i=1;i<n;i++){
int u,v;
scanf("%d%d",&u,&v);
add(u,v);
add(v,u);
}
dfs1(1,0);
dfs2(1,0,1);
for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
return 0;
}