CF600E Lomsat gelral

CF600E Lomsat gelral

• \(dsu\ on\ tree\) 入门例题.讲解.
• 注意 \(dfs\) 中各个操作的顺序,这样能保证清空轻儿子贡献时不会对重儿子造成影响.

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define mp make_pair
#define pii pair<int,int>
inline int read()
{
    int x=0;
    bool pos=1;
    char ch=getchar();
    for(;!isdigit(ch);ch=getchar())
        if(ch=='-')
            pos=0;
    for(;isdigit(ch);ch=getchar())
        x=x*10+ch-'0';
    return pos?x:-x;
}
const int MAXN=1e5+10;
vector<int> sons[MAXN];
int n;
int mxson[MAXN],siz[MAXN],col[MAXN],cnt[MAXN];
ll sum,ans[MAXN];
void getsiz(int u,int fa)
{
    siz[u]=1;
    int tot=sons[u].size();
    for(int i=0;i<tot;++i)
        {
            int v=sons[u][i];
            if(v!=fa)
                {
                    getsiz(v,u);
                    if(siz[v]>siz[mxson[u]])
                        mxson[u]=v;
                    siz[u]+=siz[v];
                }
        }
}
int mxcnt,bigson;
void add(int u,int fa,int val)
{
    cnt[col[u]]+=val;
    if(cnt[col[u]]>mxcnt)
        mxcnt=cnt[col[u]],sum=col[u];
    else if(cnt[col[u]]==mxcnt)
        sum+=col[u];
    int tot=sons[u].size();
    for(int i=0;i<tot;++i)
        {
            int v=sons[u][i];
            if(v!=fa && v!=bigson)
                add(v,u,val);
        }
}
void dfs(int u,int fa,bool keep)
{
    int tot=sons[u].size();
    for(int i=0;i<tot;++i)
        {
            int v=sons[u][i];
            if(v!=fa && v!=mxson[u])
                dfs(v,u,false);
        }
    if(mxson[u])
        dfs(mxson[u],u,true),bigson=mxson[u];
    add(u,fa,1);
    ans[u]=sum;
    bigson=0;
    if(keep==false)
        add(u,fa,-1),sum=0,mxcnt=0;
}
int main()
{
    n=read();
    for(int i=1;i<=n;++i)   
        col[i]=read();
    for(int i=1;i<n;++i)
        {
            int u=read(),v=read();
            sons[u].push_back(v);
            sons[v].push_back(u);       
        }
    getsiz(1,0);
    dfs(1,0,true);
    for(int i=1;i<=n;++i)
        printf("%I64d ",ans[i]);
    return 0;
}

转载于:https://www.cnblogs.com/jklover/p/10388875.html