Submission #1368861

#TimeUsernameProblemLanguageResultExecution timeMemory
1368861monk_is_liveRegions (IOI09_regions)C++20
100 / 100
2173 ms22124 KiB
#include <bits/stdc++.h>
using namespace std;

using ll = long long;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, R, Q;
    cin >> N >> R >> Q;

    vector<int> region(N + 1);
    vector<vector<int>> children(N + 1);

    // Read root
    cin >> region[1];

    // Read the rest
    for (int v = 2; v <= N; ++v) {
        int s;
        cin >> s >> region[v];
        children[s].push_back(v);
    }

    // Euler tour
    vector<int> tin(N + 1), tout(N + 1);
    vector<vector<int>> regionNodes(R + 1); // node ids by region
    vector<vector<int>> regionTins(R + 1);   // tin values by region

    int timer = 0;
    vector<int> it(N + 1, 0);
    vector<int> st;
    st.reserve(N);
    st.push_back(1);

    while (!st.empty()) {
        int v = st.back();

        if (it[v] == 0) {
            tin[v] = ++timer;
            regionNodes[region[v]].push_back(v);
            regionTins[region[v]].push_back(tin[v]);
        }

        if (it[v] < (int)children[v].size()) {
            int to = children[v][it[v]++];
            st.push_back(to);
        } else {
            tout[v] = timer;
            st.pop_back();
        }
    }

    // Split regions into light/heavy
    const int B = 450;
    vector<int> heavyId(R + 1, -1);
    vector<int> heavyRegions;

    for (int r = 1; r <= R; ++r) {
        if ((int)regionNodes[r].size() > B) {
            heavyId[r] = (int)heavyRegions.size();
            heavyRegions.push_back(r);
        }
    }

    // Precompute answers for heavy first region:
    // heavyAns[id_of_h][x] = answer for query (h, x)
    vector<vector<int>> heavyAns(heavyRegions.size(), vector<int>(R + 1, 0));

    for (int idx = 0; idx < (int)heavyRegions.size(); ++idx) {
        int h = heavyRegions[idx];

        // DFS with current count of h-nodes on the root-to-parent path
        vector<pair<int, int>> st2;
        st2.reserve(N);
        st2.push_back({1, 0});

        while (!st2.empty()) {
            auto [v, cnt] = st2.back();
            st2.pop_back();

            // cnt = number of ancestors of v that are in region h
            heavyAns[idx][region[v]] += cnt;

            int nextCnt = cnt + (region[v] == h ? 1 : 0);

            // Push children
            for (int i = (int)children[v].size() - 1; i >= 0; --i) {
                st2.push_back({children[v][i], nextCnt});
            }
        }
    }

    // Answer queries
    while (Q--) {
        int r1, r2;
        cin >> r1 >> r2;

        ll ans = 0;

        if (heavyId[r1] != -1) {
            // Heavy first region: O(1)
            ans = heavyAns[heavyId[r1]][r2];
        } else {
            // Light first region: iterate only over nodes in r1
            const vector<int> &vec = regionTins[r2];

            for (int u : regionNodes[r1]) {
                int L = tin[u];
                int Rr = tout[u];
                ans += upper_bound(vec.begin(), vec.end(), Rr)
                     - lower_bound(vec.begin(), vec.end(), L);
            }
        }

        cout << ans << '\n' << flush;
    }

    return 0;
}
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...