Codeforces 600E Lomsat gelral （树上启发式合并）

题目链接 Lomsat gelral

占坑……等深入理解了再来补题解……

#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b)    for (int i(a); i <= (b); ++i)

typedef long long LL;

const int N = 600010;

int n;
int cc[N], col[N], sz[N], son[N];
LL ans[N];

vector <int> v[N];

void getsize(int x, int fa){
        sz[x] = 1;
        for (auto u : v[x]){
                if (u == fa) continue;
                getsize(u, x); sz[x] += sz[u];
                if (!son[x] || sz[u] > sz[son[x]]) son[x] = u;
        }
}

bool skip[N];
LL sum = 0;
int cx = 0;
int x, y;

void edt(int x, int fa, int k){
        cc[col[x]] += k;
        if (k > 0 && cc[col[x]] >= cx){
                if (cc[col[x]] > cx)
                        sum = 0, cx = cc[col[x]];
                sum += col[x];
        }

        for (auto u : v[x]){
        if (u == fa || skip[u]) continue;
        edt(u, x, k);
    }

}

void dfs(int x, int fa, int kep = 0){
        for (auto u : v[x]){
                if (u == fa || u == son[x]) continue;
                dfs(u, x, 0);
        }

        if (son[x]) dfs(son[x], x, 1), skip[son[x]] = 1;
        edt(x, fa, 1);

        ans[x] = sum;
        if (son[x]) skip[son[x]] = 0;
        if (!kep) edt(x, fa, -1), cx = sum = 0;
}

int main(){

    scanf("%d", &n);
    rep(i, 1, n) scanf("%d", col + i);

    rep(i, 1, n - 1){
        scanf("%d%d", &x, &y);
        v[x].push_back(y);
        v[y].push_back(x);
    }

    getsize(1, 0);
    dfs(1, 0, 0);

    rep(i, 1, n) printf("%lld\n", ans[i]);
    return 0;
}