Submission #1139701

#	Time	Username	Problem	Language	Result	Execution time	Memory
1139701		selmahbn	Regions (IOI09_regions)	C++20	100 / 100	1881 ms	57272 KiB

regions

#include <bits/stdc++.h>

using namespace std;

int num = 0;
vector<int> o;

void dfs(int curr, int par, vector<int>& reg, vector<vector<int>>& orr, vector<int>& dfs_num, vector<int>& ss, vector<vector<int>>& adj) {
    dfs_num[curr] = num;
    num++;
    orr[reg[curr]].push_back(curr);
    o.push_back(curr);
    for (int neigh : adj[curr]) {
        if (neigh != par) {
            dfs(neigh, curr, reg, orr, dfs_num, ss, adj);
            ss[curr] += ss[neigh];
        }
    }
}

int main()
{
    int n, r, q;
    cin >> n >> r >> q;
    vector<int> reg(n);
    cin >> reg[0];
    reg[0]--;
    vector<vector<int>> adj(n);
    vector<int> par(n);
    par[0] = -1;
    for (int i = 1; i < n; i++) {
        cin >> par[i] >> reg[i];
        par[i]--; reg[i]--;
        adj[i].push_back(par[i]);
        adj[par[i]].push_back(i);
    }
    vector<int> dfs_num(n);
    vector<int> ss(n, 1);
    vector<vector<int>> orr(r);
    dfs(0, 0, reg, orr, dfs_num, ss, adj);
    vector<vector<int>> pre1(n);
    vector<vector<int>> pre2(n);
    int sq = sqrt(n);
    for (int i = 0; i < r; i++) {
        if ((int) orr[i].size() >= sq) {
            pre1[i].assign(r, 0);
            vector<int> dp1(n, 0);
            for (int j = 0; j < n; j++) {
                if (reg[o[j]] == i) dp1[j]++;
                if (j != 0) dp1[j] += dp1[dfs_num[par[o[j]]]];
                pre1[i][reg[o[j]]] += dp1[j];
            }
            pre2[i].assign(r, 0);
            vector<int> dp2(n, 0);
            for (int j = n-1; j >= 0; j--) {
                pre2[i][reg[o[j]]] += dp2[j];
                if (j != 0) {
                    if (reg[o[j]] == i) dp2[j]++;
                    dp2[dfs_num[par[o[j]]]] += dp2[j];
                }
            }
        }
    }
    for (int i = 0; i < q; i++) {
        int r1, r2;
        cin >> r1 >> r2;
        r1--; r2--;
        if ((int) orr[r1].size() >= sq) cout << pre1[r1][r2] << endl;
        else if ((int) orr[r2].size() >= sq) cout << pre2[r2][r1] << endl;
        else {
            int sum = 0;
            stack<int> s;
            int i1 = 0;
            int i2 = 0;
            while (i2 < (int) orr[r2].size()) {
                if (i1 >= (int) orr[r1].size() || dfs_num[orr[r2][i2]] < dfs_num[orr[r1][i1]]) {
                    while (!s.empty() && dfs_num[orr[r2][i2]] >= dfs_num[s.top()] + ss[s.top()]) s.pop();
                    sum += (int) s.size();
                    i2++;
                } else {
                    while (!s.empty() && dfs_num[orr[r1][i1]] >= dfs_num[s.top()] + ss[s.top()]) s.pop();
                    s.push(orr[r1][i1]);
                    i1++;
                }
            }
            cout << sum << endl;
        }
    }
    return 0;
}

• Subtask 1: 15 / 15
• Subtask 2: 85 / 85