SP10707 COT2 - Count on a tree II

题目描述

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

Solution

树上莫队

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
#include <iostream>
#include <cstdio>
#include <cmath>
#include <algorithm>
#define maxn 40010
#define maxm 100010
using namespace std;

int n, m, a[maxn];

int b[maxn];
void init_hash() {
for (int i = 1; i <= n; ++i) b[i] = a[i];
sort(b + 1, b + n + 1); int cnt = unique(b + 1, b + n + 1) - b - 1;
for (int i = 1; i <= n; ++i)
a[i] = lower_bound(b + 1, b + cnt + 1, a[i]) - b;
}

int blo;
struct Query {
int l, r, g, id;

friend bool operator < (const Query &u, const Query &v) {
if (u.l / blo == v.l / blo) return u.l / blo & 1 ? u.r < v.r : u.r > v.r;
return u.l < v.l;
}
} Q[maxm];

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 f[maxn][21], dep[maxn], in[maxn], out[maxn], bl[maxn * 2], c2;
void dfs(int u, int fa) {
in[u] = ++c2; bl[c2] = u;
for (int i = head[u]; ~i; i = e[i].next) {
int v = e[i].to; if (v == fa) continue;
dep[v] = dep[u] + 1; f[v][0] = u; dfs(v, u);
} out[u] = ++c2; bl[c2] = u;
}

void init_lca() {
for (int j = 1; j <= 20; ++j)
for (int i = 1; i <= n; ++i)
f[i][j] = f[f[i][j - 1]][j - 1];
}

int get_lca(int x, int y) {
if (dep[x] < dep[y]) swap(x, y);
for (int i = 20; ~i; --i)
if (dep[f[x][i]] >= dep[y]) x = f[x][i];
if (x == y) return x;
for (int i = 20; ~i; --i)
if (f[x][i] != f[y][i]) x = f[x][i], y = f[y][i];
return f[x][0];
}

int cnt[maxn], Ans[maxm], ans; bool use[maxn];
inline void modify(int k) {
if (!use[k] && ++cnt[a[k]] == 1) ++ans;
if (use[k] && --cnt[a[k]] == 0) --ans;
use[k] ^= 1;
}

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

cin >> n >> m; blo = sqrt(n);
for (int i = 1; i <= n; ++i) cin >> a[i];
for (int i = 1; i < n; ++i) {
int x, y; cin >> x >> y;
add_edge(x, y); add_edge(y, x);
}
init_hash(); dep[1] = 1; dfs(1, 0); init_lca();
for (int i = 1; i <= m; ++i) {
int x, y; cin >> x >> y;
if (in[x] > in[y]) swap(x, y);
if (in[x] <= in[y] && out[y] <= out[x])
Q[i] = { in[x], in[y], 0, i };
else Q[i] = { out[x], in[y], in[get_lca(x, y)], i };
} sort(Q + 1, Q + m + 1);
int l = Q[1].l, r = l - 1;
for (int i = 1; i <= m; ++i) {
while (r < Q[i].r) modify(bl[++r]);
while (l > Q[i].l) modify(bl[--l]);
while (r > Q[i].r) modify(bl[r--]);
while (l < Q[i].l) modify(bl[l++]);
if (Q[i].g) modify(bl[Q[i].g]);
Ans[Q[i].id] = ans;
if (Q[i].g) modify(bl[Q[i].g]);
}
for (int i = 1; i <= m; ++i) cout << Ans[i] << "\n";
return 0;
}