Submission #763781

#TimeUsernameProblemLanguageResultExecution timeMemory
763781izanbfRegions (IOI09_regions)C++14
95 / 100
3293 ms131072 KiB
#include <bits/stdc++.h>
using namespace std;

using vi = vector<int>;
using vvi = vector<vi>;

const int BIG = 250;
vi r, tin, tout, big_id;
vvi adj, tins, touts, by_r, dp;
int timer;

void euler(int u, int p) {
    tin[u] = timer++;
    tins[r[u]].push_back(tin[u]);
    for (int v : adj[u]) {
        if (v != p) euler(v, u);
    }
    tout[u] = timer++;
    touts[r[u]].push_back(tout[u]);
}

int count_in_range(const vi& v, int l, int r) {
    if (v.empty()) return 0;
    int ri = upper_bound(v.begin(), v.end(), r) - v.begin();
    int li = upper_bound(v.begin(), v.end(), l) - v.begin();
    return ri - li;
}

void precompute(int rg, int p, int u) {
    int rid = big_id[rg];
    dp[rid][u] = (r[u] == rg);
    for (int v : adj[u]) {
        if (v != p) {
            precompute(rg, u, v);
            dp[rid][u] += dp[rid][v];
        }
    }
}

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    
    int N, R, Q;
    cin >> N >> R >> Q;

    r = tin = tout = vi(N);
    big_id = vi(R);
    adj = vvi(N);
    tins = touts = by_r = vvi(R);

    vi p(N, -1);

    cin >> r[0]; --r[0];
    by_r[r[0]].push_back(0);

    for (int i = 1; i < N; ++i) {
        cin >> p[i] >> r[i]; --p[i], --r[i];
        by_r[r[i]].push_back(i);
        adj[i].push_back(p[i]);
        adj[p[i]].push_back(i);
    }

    int bigs = 0;
    for (int rg = 0; rg < R; ++rg) {
        if (by_r[rg].size() >= BIG) big_id[rg] = bigs++;
    }

    dp = vvi(bigs, vi(N)); // O(N/BIG N)
    for (int rg = 0; rg < R; ++rg) {
        if (by_r[rg].size() >= BIG) precompute(rg, 0, 0);
    }

    timer = 1;
    euler(0, 0);

    map<pair<int,int>,long long> mem;

    while (Q--) {
        int r1, r2;
        cin >> r1 >> r2; --r1, --r2;
        if (by_r[r1].size() < BIG and by_r[r2].size() >= BIG) { // r1 small r2 big  O(Q BIG)
            long long cnt = 0;
            int rid = big_id[r2];
            for (int u : by_r[r1]) {
                cnt += dp[rid][u];
            }
            cout << cnt << endl;
        }
        else { //O(Q BIG log N)
            if (not mem.count({r1,r2})) {
                long long cnt = 0;
                for (int u : by_r[r2]) {
                    cnt += count_in_range(tins[r1], 0, tin[u])
                         - count_in_range(touts[r1], 0, tin[u]);
                }
                mem[{r1,r2}] = cnt;
            }
            cout << mem[{r1,r2}] << endl;
        }
    }
}




#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...