600E - Lomsat gelral(找子树多颜色问题)（入门）

Posted 2021-03-13 starve

题：https://codeforces.com/problemset/problem/600/E

题意：一棵树有n个结点，每个结点都是一种颜色，每个颜色有一个编号，求树中每个子树的最多的颜色编号的和，对于每个结点都输出答案。

分析：考虑暴力算法，对于每个节点只是清空计数数组，再对其子树颜色进行统计，复杂度o(n^2);

　　　接着我们发现最后一个子树的清空是必要的，所以我们把重儿子当作这个子树，就可以让复杂度降为o（nlogn）级别；

#include<bits/stdc++.h>
using namespace std;
#define pb push_back
typedef long long ll;
const int M=2e5+5;
ll ans[M];
ll countt[M],sz[M],son[M],col[M],vis[M];
vector<int>g[M];
ll nownum,maxx;
void dfs1(int u,int fa){
    sz[u]=1;
    for(int i=0;i<g[u].size();i++){
        int v=g[u][i];
        if(v!=fa){
            dfs1(v,u);
            sz[u]+=sz[v];
            if(sz[v]>sz[son[u]])
                son[u]=v;
        }
    }
}

void update(int u,int k,int fa){//k的取值为1和-1,分别对应累加和清除
    countt[col[u]]+=k;
    if(maxx<countt[col[u]])
        maxx=countt[col[u]],nownum=col[u];
    else if(maxx==countt[col[u]])
        nownum+=col[u];
    for(int i=0;i<g[u].size();i++){
        int v=g[u][i];
        if(v!=fa&&!vis[v])
            update(v,k,u);
    }
}
void dfs2(int u,int fa,int sign){
    //cout<<u<<endl;
    for(int i=0;i<g[u].size();i++){
        int v=g[u][i];
        if(v!=fa&&son[u]!=v)
            dfs2(v,u,0);
    }
    if(son[u])//再访问重儿子,一定要打上vis标记
        dfs2(son[u],u,1),vis[son[u]]=1;
    update(u,1,fa);
    ans[u]=nownum;
    if(son[u])
        vis[son[u]]=0;
    if(!sign)///轻儿子就消除影响 
        update(u,-1,fa),nownum=maxx=0;
    
}
int main(){
    int n;
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
        scanf("%d",&col[i]);
    for(int u,v,i=1;i<n;i++){
        scanf("%d%d",&u,&v);
        g[u].pb(v);
        g[v].pb(u);
    }
    dfs1(1,0);
    
    dfs2(1,0,0);
    for(int i=1;i<=n;i++)
        printf("%I64d ",ans[i]);
    return 0;
}

View Code