Luogu P5901 [IOI2009]regions

题目描述

https://www.luogu.com.cn/problem/P5901

简要题意：给定一棵 $n$ 个点有根树，树上每个点有一个颜色 $col_i$，现在有 $m$ 个询问，每次询问给定两个颜色 $(x,y)$，问有多少对节点 $(u,v)$ 满足 $u$ 是 $v$ 的祖先且 $u$ 的颜色是 $x$，$v$ 的颜色是 $y$

$n\le 2\times 10^5,m\le 2\times 10^5$

Solution

注意到询问满足对二元组的自然根号

那么我们现在可以枚举颜色数较小的那种颜色，如果其为 $u$，那么相当于查询子树中某种颜色的个数，如果其为 $v$，相当于查询树链上的某种颜色的个数，两者皆可在 $dfs$ 的时候处理

时间复杂度 $O(n\sqrt n)$

#include <iostream>
#include <vector>
#include <map>
#define maxn 200010
using namespace std;

int n, m, r, a[maxn], num[maxn];

struct Edge {
    int to, next;
} e[maxn * 2]; int c1, head[maxn];
inline void add_edge(int u, int v) {
    e[c1].to = v; e[c1].next = head[u]; head[u] = c1++;
}

int ans[maxn], Q[maxn];
map<pair<int, int>, int> mp;

vector<pair<int, int>> A[maxn], B[maxn];
int cnt[maxn];
void dfs1(int u, int fa) {
    for (auto [k, id] : A[a[u]]) ans[id] += cnt[k];
    ++cnt[a[u]];
    for (int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to; if (v == fa) continue;
        dfs1(v, u); 
    } --cnt[a[u]];
}

void dfs2(int u, int fa) {
    for (auto [k, id] : B[a[u]]) ans[id] -= cnt[k];
    ++cnt[a[u]];
    for (int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to; if (v == fa) continue;
        dfs2(v, u); 
    }
    for (auto [k, id] : B[a[u]]) ans[id] += cnt[k];
}

int main() { fill(head, head + maxn, -1);
    ios::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);

    cin >> n >> r >> m >> a[1]; ++num[a[1]];
    for (int i = 2; i <= n; ++i) {
        int x; cin >> x >> a[i];
        add_edge(x, i); ++num[a[i]];
    }
    for (int i = 1; i <= m; ++i) {
        int x, y; cin >> x >> y;
        if (mp.find({ x, y }) != mp.end()) { Q[i] = mp[{ x, y }]; continue; }
        Q[i] = mp[{ x, y }] = i; 
        if (num[x] < num[y] || num[x] == num[y] && x < y) B[x].emplace_back(y, i);
        else A[y].emplace_back(x, i); 
    } dfs1(1, 0); dfs2(1, 0);
    for (int i = 1; i <= m; ++i) cout << ans[Q[i]] << "\n";
    return 0;
}