Educational Codeforces Round 2 E. Lomsat gelral

問題

codeforces.com

解法

"データ構造をマージする一般的なテク" を使います。なにそれ, という人は以下のページをご覧ください。
topcoder.g.hatena.ne.jp

今回の問題の場合, M[v][color] = (頂点 v 以下の部分木で, 色が color の頂点がいくつあるか) というのを, 正直に更新していくと TLE しそうです。ですが, 頂点 v と頂点 u の情報をマージするときに, M[v].size() と M[u].size() のうち, 大きいサイズの方に情報をマージするようにすると, 各色の情報の移動回数がたかだか log N 回になるので, 全体計算量は O(N log N) で収まります。

他に, maxi[v] = (v 以下の部分木で支配的な色の個数), dSum[v] = (v 以下の部分木で支配的な色の数の和(C[0] + C[4] + ... みたいな)) というのを用意しておくと, 更新処理が高速に出来ます。

const int MAXN = 100100;
int C[MAXN], maxi[MAXN];
vector<int> G[MAXN];
ll ans[MAXN], dSum[MAXN];
map<int, int> M[MAXN];

void merge(int v, int u) {
    if (M[v].size() < M[u].size()) {
        swap(M[v], M[u]);
        swap(dSum[v], dSum[u]);
        swap(maxi[v], maxi[u]);
    }
    for (auto p : M[u]) {
        M[v][p.first] += p.second;
        if (maxi[v] < M[v][p.first]) {
            dSum[v] = 0;
            maxi[v] = M[v][p.first];
        }
        if (M[v][p.first] == maxi[v]) {
            dSum[v] += p.first;
        }
    }
}

void dfs(int v, int p) {
    for (int ch : G[v]) {
        if (ch == p) continue;
        dfs(ch, v);
        merge(v, ch);
    }
    M[v][C[v]] += 1;
    if (M[v][C[v]] > maxi[v]) {
        dSum[v] = 0;
        maxi[v] = M[v][C[v]];
    }
    if (M[v][C[v]] == maxi[v]) dSum[v] += C[v];
    ans[v] = dSum[v];
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int n;
    cin >> n;
    for (int i = 0; i < n; i++) cin >> C[i];
    for (int i = 0; i < n-1; i++) {
        int x, y;
        cin >> x >> y;
        x--; y--;
        G[x].push_back(y);
        G[y].push_back(x);
    }
    dfs(0, -1);
    for (int i = 0; i < n; i++) {
        cout << ans[i] << " ";
    }
    cout << endl;
    return 0;
}