Codeforces 600E. Lomsat gelral(树上启发式合并)

题意：

n个点的有根树，以1为根，每个点有一种颜色。我们称一种颜色占领了一个子树当且仅当没有其他颜色在这个子树中出现得比它多。求占领每个子树的所有颜色之和。

题解：

树上启发式合并, [Tutorial] Sack (dsu on tree) - Codeforces

#include<bits/stdc++.h>
using namespace std;
#define N 100005
#define LL long long
int n, Max, cnt[N], sz[N], col[N], big[N];
LL sum, ans[N];
vector<int> g[N];
void add(int u, int fa, int x)
{
    cnt[col[u]]+=x;
    if(cnt[col[u]]>Max) sum=col[u], Max=cnt[col[u]];
    else if(cnt[col[u]]==Max) sum+=col[u];
    for(auto v:g[u])
        if(v!=fa&&!big[v])
            add(v, u, x);
}
void dfs(int u, int fa, bool keep)
{
    int mx=-1, bigchild=-1;
    for(auto v:g[u])
        if(v!=fa&&sz[v]>mx)
            mx=sz[v], bigchild=v;
    for(auto v:g[u])
        if(v!=fa&&v!=bigchild)
            dfs(v, u, 0);
    if(bigchild!=-1)
        dfs(bigchild, u, 1), big[bigchild]=1;
    add(u, fa, 1);
    if(bigchild!=-1)
        big[bigchild]=0;
    ans[u]=sum;
    if(keep==0)
        add(u, fa, -1), Max=sum=0;
}

void getsz(int u, int fa)
{
    sz[u]=1;
    for(auto v:g[u])
        if(v!=fa)
        {
            getsz(v, u);
            sz[u]+=sz[v];
        }
}
void solve()
{
    cin>>n;
    int u, v;
    for(int i=1; i<=n; i++) cin>>col[i];
    for(int i=1; i<n; i++) cin>>u>>v, g[u].push_back(v), g[v].push_back(u);
    getsz(1, 0);
    dfs(1, 0, 1);
    for(int i=1; i<=n; i++) cout << ans[i] << " ";
}
int main()
{
    solve();
    return 0;
}