CS Academy Uniform Trees

リンク : https://csacademy.com/contest/archive/task/uniform-trees/statement/

問題概要

N頂点の木が与えられる、それそれの頂点にラベルがつけられる。以下の条件を満たす頂点集合の通り数を求めよ。
・頂点集合のうちある１つの頂点はほかのすべての頂点の祖先となる。これを根とする。
・根以外の頂点は頂点集合中の最も近い祖先とつなぐ。

解説

dp[i][j]を頂点iの子孫の頂点集合に関して、つながれた祖先からの頂点がjのラベルをつけられた働きをするときの通り数とすると、dp[i][j]=∏{(​x∈child(i))dp[x][j]+1}-1となる。このdpはi<=2×、j<=2×、となるため、静的に確保してしまうと、MLEする。そこでjに関しては中身が0のことが多いため、mapでもっておくことを考える。すると、子孫からのdpをうけとるときにマージテクなるものをつかうことができる。これは大きいほうに小さいものを入れることを繰り返せばからNlogNに落とせるといったテクニックである。
メモ:マージテクは計算量空間量ともに削減する。

LL n, ans, a[210000],dp[210000];
map<LL, LL>mp[210000];
vector<vector<LL>>g;
void merge(LL a, LL b) {
    if (mp[a].size() < mp[b].size()) {
        swap(mp[a], mp[b]);
        swap(dp[a], dp[b]);
    }
    for (auto it = mp[b].begin(); it != mp[b].end(); it++) {
        if (mp[a].count(it->first)) {
            (dp[a] += (mp[a][it->first] *( it->second - 1 + MOD)) % MOD) %= MOD;
            mp[a][it->first] = (mp[a][it->first] * it->second) % MOD;
        }
        else {
            (dp[a] += it->second)%=MOD;
            mp[a][it->first] = it->second;
        }
    }
}
void dfs(LL cur, LL par) {
    rep(i, g[cur].size()) {
        LL v = g[cur][i];
        if (v == par)continue;
        dfs(v, cur);
        merge(cur, v);
    }
    LL num = (1 + dp[cur] + MOD - mp[cur].size()) % MOD;
    (ans += num) %= MOD;
    if (mp[cur].count(a[cur])) {
        (dp[cur] += num) %= MOD;
        (mp[cur][a[cur]] += num) %= MOD;
    }
    else {
        (dp[cur] += num+1) %= MOD;
        (mp[cur][a[cur]] += num+1) %= MOD;
    }
}
int main() {
    cin >> n;
    g.resize(n);
    rep(i, n) {
        LL x, y;
        cin >> x >> a[i];
        x--;
        if (i) {
            g[x].push_back(i);
            g[i].push_back(x);
        }
    }
    dfs(0, -1);
    cout << ans << endl;
    return 0;
}