Codeforces 600 E - Lomsat gelral

E - Lomsat gelral

思路：

树上启发式合并

代码：

#include<bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define pi acos(-1.0)
#define LL long long
//#define mp make_pair
#define pb push_back
#define ls rt<<1, l, m
#define rs rt<<1|1, m+1, r
#define ULL unsigned LL
#define pll pair<LL, LL>
#define pii pair<int, int>
#define piii pair<int,pii>
#define mem(a, b) memset(a, b, sizeof(a))
#define fio ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define fopen freopen("in.txt", "r", stdin);freopen("out.txt", "w", stout);
//head

const int N = 1e5 + 5;
int color[N], sz[N], bs[N], mx = 0;
LL ans[N], t = 0;
vector<int>g[N];
map<int, int>mp;
void dfs(int o, int u) {
    sz[u] = 1;
    for (int v : g[u]) {
        if(v != o) {
            dfs(u, v);
            sz[u] += sz[v];
            if(sz[v] > sz[bs[u]]) bs[u] = v;
        }
    }
}
void ADD(int o, int u) {
    mp[color[u]] ++;
    if(mp[color[u]] > mx) t = color[u], mx = mp[color[u]];
    else if(mp[color[u]] == mx) t += color[u];
    for (int v : g[u]) {
        if(v != o) {
            ADD(u, v);
        }
    }
}
void DFS(int o, int u) {
    for (int v : g[u]) {
        if(v != o && v != bs[u]) {
            DFS(u, v);
            mp.clear();
            mx = 0;
            t = 0;
        }
    }
    if(bs[u]) DFS(u, bs[u]);
    for (int v : g[u]) {
        if(v != o && v != bs[u]) {
            ADD(u, v);
        }
    }
    mp[color[u]] ++;
    if(mp[color[u]] > mx) t = color[u], mx = mp[color[u]];
    else if(mp[color[u]] == mx) t += color[u];
    ans[u] = t;
}
int main() {
    int n, u, v;
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) scanf("%d", &color[i]);
    for (int i = 1; i < n; i++) {
        scanf("%d %d", &u, &v);
        g[u].pb(v);
        g[v].pb(u);
    }
    dfs(0, 1);
    DFS(0, 1);
    for (int i = 1; i <= n; i++) printf("%lld%c", ans[i], " \n"[i==n]);
    return 0;
}