Educational Codeforces Round 2 E: Lomsat gelral

問題概要

原文 → Problem - E - Codeforces

頂点からなる根付き木が与えられる。根は頂点 であり、木の各頂点には色 がついている。

各頂点 について、 を根とする部分木に存在する頂点の中で、塗られている個数が最も多い色を求め (タイもすべて数えるため複数ある場合がある)、色 id の合計を出力せよ。

解説

データ構造をマージする一般的なテク、通称「マージテク」を使って解くことを考えます。

現在の頂点 の情報と子 の情報をマージする際に、必ず「サイズが小さい方を大きい方に向けてマージする」ようにします。こうすることでマージ全体にかかる時間がならし 時間になります。map の計算量も含めて、 で処理可能です。

ソースコード

using Graph = vector< vector<int> >;

int N, col[100010], res[100010];
map<int, int> exist[100010];
map<int, int> sum[100010];

void solve(Graph &G, int cur, int par=-1) {
    exist[cur][ col[cur] ] = 1;
    sum[cur][1] = col[cur];
    for(auto to : G[cur]) {
        if(to == par) continue;
        solve(G, to, cur);

        // サイズを比較して、小さい方を大きい方に向けてマージ
        if(exist[cur].size() < exist[to].size()) {
            swap(exist[cur], exist[to]);
            swap(sum[cur], sum[to]);
        }

        // to を cur にマージ
        for(auto e : exist[to]) {
            int id, cnt; tie(id, cnt) = e;
            if(exist[cur].count(id)) {
                int cnt_ = exist[cur][id];
                cnt += cnt_;
                exist[cur][id] = cnt;
                sum[cur][cnt_] -= id;
                sum[cur][cnt ] += id;
            }
            else {
                exist[cur][id] = cnt;
                sum[cur][cnt] += id;
            }
        }
    }
    int ans = (--sum[cur].end())->second;
    res[cur] = ans;
}

signed main() {
    scanf("%lld", &N);
    for(int i=0; i<N; i++) {
        scanf("%lld", &col[i]);
    }

    Graph G(N);
    for(int i=0; i<N-1; i++) {
        int u, v;
        scanf("%lld%lld", &u, &v);
        u--; v--;
        G[u].push_back(v);
        G[v].push_back(u);
    }

    solve(G, 0);
    for(int i=0; i<N; i++) {
        printf("%lld%c", res[i], " \n"[i+1 == N]);
    }
    return 0;
}