Regions

Solution

Sqrt Decomposition

Let $B = \sqrt{N}$. A region is large if it has $\ge B$ nodes, small otherwise. There are at most $\sqrt{N}$ large regions.

\textbf{Case 1: $r_1$ is large.} Precompute via DFS: maintain a counter of ancestors in $r_1$. At each node $u$ of region $r_2$, add the counter value to $\text{ans}[r_2]$. Total precomputation: $O(N \sqrt{N})$.

\textbf{Case 2: $r_2$ is large.} Precompute $\text{subtreeR2}[u]$ = number of descendants of $u$ in $r_2$ (post-order traversal). For each node $u$ in region $r_1$, add $\text{subtreeR2}[u]$ (excluding $u$ itself if $u \in r_2$, but since $r_1 \ne r_2$ or a node is in one region, this is automatic).

Case 3: Both small. For each node $u$ in $r_1$, binary-search on the sorted Euler-tour times of $r_2$ to count how many $r_2$-nodes have $\text{tin}$ in $[\text{tin}[u]+1, \text{tout}[u]]$. Total work: $O(|r_1| \cdot \log |r_2|)$.

Complexity

• Precomputation: $O(N\sqrt{N})$.

• Per query: $O(1)$ for large-region lookups; $O(|r_1| \log |r_2|)$ for small-small (bounded by $O(B \log B) = O(\sqrt{N} \log N)$).

• Space: $O(N\sqrt{N})$.

C++ Solution

#include <bits/stdc++.h>
using namespace std;

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);
    vector<vector<int>> regionNodes(R + 1);

    cin >> region[1];
    regionNodes[region[1]].push_back(1);

    for(int i = 2; i <= N; i++){
        int p;
        cin >> p >> region[i];
        children[p].push_back(i);
        regionNodes[region[i]].push_back(i);
    }

    // Euler tour (iterative DFS)
    vector<int> tin(N + 1), tout(N + 1), euler;
    int timer = 0;
    {
        stack<pair<int,bool>> stk;
        stk.push({1, false});
        while(!stk.empty()){
            auto [u, leaving] = stk.top(); stk.pop();
            if(leaving){ tout[u] = timer - 1; continue; }
            tin[u] = timer++;
            euler.push_back(u);
            stk.push({u, true});
            for(int i = (int)children[u].size()-1; i >= 0; i--)
                stk.push({children[u][i], false});
        }
    }

    int B = max(1, (int)sqrt(N));

    // Sort region nodes by tin
    for(int r = 1; r <= R; r++)
        sort(regionNodes[r].begin(), regionNodes[r].end(),
             [&](int a, int b){ return tin[a] < tin[b]; });

    // Identify large regions
    vector<int> largeId(R + 1, -1);
    vector<int> largeRegions;
    for(int r = 1; r <= R; r++)
        if((int)regionNodes[r].size() >= B){
            largeId[r] = largeRegions.size();
            largeRegions.push_back(r);
        }
    int nLarge = largeRegions.size();

    // Precompute: large r1 -> all r2
    vector<vector<long long>> ansLargeR1(nLarge, vector<long long>(R + 1, 0));
    for(int li = 0; li < nLarge; li++){
        int r1 = largeRegions[li];
        int curCount = 0;
        // DFS maintaining ancestor count in r1
        stack<pair<int,bool>> stk;
        stk.push({1, false});
        while(!stk.empty()){
            auto [u, leaving] = stk.top(); stk.pop();
            if(leaving){
                if(region[u] == r1) curCount--;
                continue;
            }
            if(region[u] == r1) curCount++;
            // u has curCount ancestors in r1 (including itself if in r1)
            // We want strict ancestors, so subtract 1 if u is in r1
            ansLargeR1[li][region[u]] += curCount - (region[u] == r1 ? 1 : 0);
            stk.push({u, true});
            for(int i = (int)children[u].size()-1; i >= 0; i--)
                stk.push({children[u][i], false});
        }
    }

    // Precompute: all r1 -> large r2
    vector<vector<long long>> ansLargeR2(R + 1, vector<long long>(nLarge, 0));
    for(int li = 0; li < nLarge; li++){
        int r2 = largeRegions[li];
        // Compute subtreeR2[u] using reverse euler order
        vector<int> sub(N + 1, 0);
        for(int i = (int)euler.size()-1; i >= 0; i--){
            int u = euler[i];
            sub[u] = (region[u] == r2) ? 1 : 0;
            for(int c : children[u]) sub[u] += sub[c];
        }
        for(int u = 1; u <= N; u++){
            int cnt = sub[u] - (region[u] == r2 ? 1 : 0);
            if(cnt > 0) ansLargeR2[region[u]][li] += cnt;
        }
    }

    // Count nodes of region r with tin in [lo, hi]
    auto countInRange = [&](int r, int lo, int hi) -> int {
        auto& v = regionNodes[r];
        auto it_lo = lower_bound(v.begin(), v.end(), lo,
            [&](int node, int val){ return tin[node] < val; });
        auto it_hi = upper_bound(v.begin(), v.end(), hi,
            [&](int val, int node){ return val < tin[node]; });
        return (int)(it_hi - it_lo);
    };

    while(Q--){
        int r1, r2;
        cin >> r1 >> r2;
        long long ans = 0;

        if(largeId[r1] >= 0){
            ans = ansLargeR1[largeId[r1]][r2];
        } else if(largeId[r2] >= 0){
            ans = ansLargeR2[r1][largeId[r2]];
        } else {
            for(int u : regionNodes[r1])
                ans += countInRange(r2, tin[u] + 1, tout[u]);
        }
        cout << ans << "\n";
        cout.flush();
    }
    return 0;
}

Notes

The sqrt decomposition is essential for handling online queries efficiently. The three cases (large $r_1$, large $r_2$, both small) each exploit different structural properties. The Euler-tour representation converts the ancestor-descendant relationship to a range query, enabling binary search on sorted entry times.