luogu CF600E Lomsat gelral |dsu on tree

一棵树有n个结点,每个结点都是一种颜色,每个颜色有一个编号,求树中每个子树的最多的颜色编号的和。


启发式搜索

#include<vector>
#include<cstdio>
#include<iostream>
#include<algorithm>
using namespace std;
#define ll long long
const int MAXN=1e5+10;
inline int read() {
    char c = getchar(); int x = 0, f = 1;
    while(c < '0' || c > '9') {if(c == '-') f = -1; c = getchar();}
    while(c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
    return x * f;
}
int N,col[MAXN],son[MAXN],siz[MAXN],cnt[MAXN],Mx,Son;
ll sum=0,ans[MAXN];
vector<int>v[MAXN];
inline void dfs(int x,int fa){//预处理重边
    siz[x]=1;
    for(int i=0;i<v[x].size();i++){
        int to=v[x][i];
        if(to==fa)continue;
        dfs(to,x);
        siz[x]+=siz[to];
        if(siz[to]>siz[son[x]])son[x]=to;
    }
}
inline void add(int x,int fa,int val){
    cnt[col[x]]+=val;
    if(cnt[col[x]]>Mx)Mx=cnt[col[x]],sum=col[x];
    else if(cnt[col[x]]==Mx)sum+=(ll)col[x];
    for(int i=0;i<v[x].size();i++){
        int to=v[x][i];
        if(to==fa||to==Son)continue;
        add(to,x,val);
    }
}
inline void dfs2(int x,int fa,int opt){
    for(int i=0;i<v[x].size();i++){
        int to=v[x][i];
        if(to==fa||to==son[x])continue;
        dfs2(to,x,0);//轻边
    }
    if(son[x])dfs2(son[x],x,1),Son=son[x];
    add(x,fa,1); Son=0;
    ans[x]=sum;
    if(!opt)add(x,fa,-1),sum=0,Mx=0;
}
signed main(){
    N=read();
    for(int i=1;i<=N;i++)col[i]=read();
    for(int i=1;i<N;i++){
        int x=read(),y=read();
        v[x].push_back(y);
        v[y].push_back(x);
    }
    dfs(1,0);
    dfs2(1,0,0);
    for(int i=1;i<=N;i++)
    printf("%lld ",ans[i]);
    return 0;
}

猜你喜欢

转载自www.cnblogs.com/naruto-mzx/p/12111578.html