codeforces 600 E. Lomsat gelral(dsu on tree)

題目鏈接：https://codeforces.com/contest/600/problem/E

思路：dsu on tree模板

#include <cstdio>
#include <cstdlib>
#include <cassert>
#include <cstring>
#include <bitset>
#include <cmath>
#include <cctype>
#include <iostream>
#include <algorithm>
#include <string>
#include <vector>
#include <queue>
#include <map>
#include <set>
#include <sstream>
#include <iomanip>
using namespace std;
typedef long long ll;
typedef vector<long long> VI;
typedef unsigned long long ull;
const ll inff = 0x3f3f3f3f3f3f3f3f;
#define FOR(i,a,b) for(int i(a);i<=(b);++i)
#define FOL(i,a,b) for(int i(a);i>=(b);--i)
#define SZ(x) ((long long)(x).size())
#define REW(a,b) memset(a,b,sizeof(a))
#define inf int(0x3f3f3f3f)
#define si(a) scanf("%d",&a)
#define sl(a) scanf("%lld",&a)
#define sd(a) scanf("%lf",&a)
#define ss(a) scanf("%s",a)
#define mod ll(998244353)
#define pb push_back
#define eps 1e-6
#define lc d<<1
#define rc d<<1|1
#define Pll pair<ll,ll>
#define P pair<int,int>
#define pi acos(-1)
ll tot,sz[100008],son[100008],head[100008],n,x,y,c,Max,a[100008],vis,qw[100008],ans[100008];
struct node{
int next,to;}e[300008];
void add(int a,int b)
{
    e[tot].next=head[a];
    e[tot].to=b;
    head[a]=tot++;
}
void dfs(int u,int f)
{
    sz[u]=1;int maxson=0;
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v=e[i].to;
        if(v==f) continue;
        dfs(v,u);sz[u]+=sz[v];
        if(sz[v]>maxson) son[u]=v,maxson=sz[v];
    }
}
void cal(int u,int fa,int z)
{
    qw[a[u]]+=z;
    if(z>0&&qw[a[u]]==Max) c+=a[u];
    else if(z>0&&qw[a[u]]>Max) c=a[u],Max=qw[a[u]];
    for(int i=head[u];i!=-1;i=e[i].next)
     if(e[i].to!=fa&&e[i].to!=vis) cal(e[i].to,u,z);
}
void dsu(int u,int T,int fa)//x子樹根節點，T表示x是否是重兒子
{
    for(int i=head[u];i!=-1;i=e[i].next)//遍歷輕兒子
     if(e[i].to!=fa&&e[i].to!=son[u]) dsu(e[i].to,0,u);
    if(son[u]) dsu(son[u],1,u),vis=son[u];//遍歷重兒子
    cal(u,fa,1);ans[u]=c;//計算整棵子樹
    if(son[u]) vis=0;
    if(!T) cal(u,fa,-1),c=Max=0;//如果不是重兒子，那麼擦除
}
int main()
{
    cin.tie(0);
    cout.tie(0);
    cin>>n;REW(head,-1);
    FOR(i,1,n) sl(a[i]);
    FOR(i,1,n-1)
    {
        sl(x),sl(y);
        add(x,y);
        add(y,x);
    }
    dfs(1,0);dsu(1,0,0);
    FOR(i,1,n) printf("%lld ",ans[i]);
    puts("");
    return 0;
}