codeforces600E. Lomsat gelral(dsu on tree)

dsu on tree
先分轻重儿子
先处理轻边,再处理重儿子
再加上轻儿子的答案

#include<iostream>
#include<cstdio>
#include<queue>
#include<algorithm>
#include<cmath>
#include<ctime>
#include<set>
#include<map>
#include<stack>
#include<cstring>
#pragma GCC optimize(2)
#define inf 2147483647
#define ls rt<<1
#define rs rt<<1|1
#define lson ls,nl,mid,l,r
#define rson rs,mid+1,nr,l,r
#define N 1000010
#define For(i,a,b) for(long long i=a;i<=b;++i)
#define p(a) putchar(a)
#define g() getchar()

using namespace std;
long long n,col[N],son[N],size[N],cnt[N],Max,Son;
long long sum,ans[N],x,y;

struct node{
    long long n;
    node *next;
}*e[N];

void in(long long &x){
    long long y=1;char c=g();x=0;
    while(c<'0'||c>'9'){if(c=='-')y=-1;c=g();}
    while(c<='9'&&c>='0'){ x=(x<<1)+(x<<3)+c-'0';c=g();}
    x*=y;
}
void o(long long x){
    if(x<0){p('-');x=-x;}
    if(x>9)o(x/10);
    p(x%10+'0');
}

void push(long long x,long long y){
    node *p;
    p=new node();
    p->n=y;
    if(e[x]==0)
        e[x]=p;
    else{
        p->next=e[x]->next;
        e[x]->next=p;
    }
}

void dfs(long long x,long long fa){//重儿子
    size[x]=1;
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa) continue;
        dfs(i->n,x);
        size[x]+=size[i->n];
        if(size[i->n]>size[son[x]])
            son[x]=i->n;
    }
}

void add(long long x,long long fa,long long val){
    cnt[col[x]]+=val;
    if(cnt[col[x]]>Max){
        Max=cnt[col[x]];
        sum=col[x];
    }
    else
        if(cnt[col[x]]==Max)
            sum+=col[x];
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa||i->n==Son) continue;
        add(i->n,x,val);
    }
}

void dfs2(long long x,long long fa,long long opt){
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa) continue;
        if(i->n!=son[x]) dfs2(i->n,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;
        Max=0;
    }
}

int main(){
    in(n);
    For(i,1,n)
        in(col[i]);
    For(i,1,n-1){
        in(x);in(y);
        push(x,y);
        push(y,x);
    }
    dfs(1,0);
    dfs2(1,1,1);
    For(i,1,n)
        o(ans[i]),p(' ');
    return 0;
}

 

posted @ 2019-07-25 10:47  WeiAR  阅读(137)  评论(0编辑  收藏  举报