U41492 树上数颜色

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
#include<bits/stdc++.h>
using namespace std;
const int MAXn = 1e5;
const int MAXm = MAXn;

template <typename T>
inline void read(T &a) {
char c;for (c = getchar(); (c < '0' || c > '9') && c != '-'; c = getchar());bool f = c == '-';T x = f ? 0 : (c ^ '0');for (c = getchar(); c >= '0' && c <= '9'; c = getchar()) {x = x * 10 + (c ^ '0');}a = f ? -x : x;
}
template <typename T, typename ...Argv>
inline void read(T &a, Argv &...argv) {
read(a), read(argv...);
}

int head[MAXn + 10], cntnex, nex[MAXm * 2 + 10], to[MAXm * 2 + 10];
inline void Insert(int u, int v) {
nex[++cntnex] = head[u];
head[u] = cntnex;
to[cntnex] = v;
}

int fa[MAXn + 10], hson[MAXn + 10], siz[MAXn + 10];
void Dfs1(int cur, int f) {
fa[cur] = f;
siz[cur] = 1;
int mx = 0;
for (int i = head[cur]; i; i = nex[i]) {
if (to[i] == fa[cur]) continue;
Dfs1(to[i], cur);
siz[cur] += siz[to[i]];
if (mx < siz[to[i]]) {
mx = siz[to[i]];
hson[cur] = to[i];
}
}
}
int color[MAXn + 10], ans[MAXn + 10];
int cntbuc, buc[MAXn + 10];
void Dfs3(int cur, int hson) {
if (!buc[color[cur]]) ++cntbuc;
++buc[color[cur]];
for (int i = head[cur]; i; i = nex[i]) {
if (to[i] == fa[cur] || to[i] == hson) continue;
Dfs3(to[i], hson);
}
}
void Dfs2(int cur) {
if (!hson[cur]) {
ans[cur] = 1;
buc[color[cur]] = 1;
cntbuc = 1;
} else {
for (int i = head[cur]; i; i = nex[i]) {
if (to[i] == fa[cur] || to[i] == hson[cur]) continue;
Dfs2(to[i]);
memset(buc, 0, sizeof(buc));
cntbuc = 0;
}
Dfs2(hson[cur]);
Dfs3(cur, hson[cur]);
ans[cur] = cntbuc;
}
}

int n, q;
signed main() {
read(n);
for (int i = 1, u, v; i < n; ++i) {
read(u, v);
Insert(u, v); Insert(v, u);
}
for (int i = 1; i <= n; ++i) {
read(color[i]);
}
Dfs1(1, 0);
Dfs2(1);
read(q);
for (int i = 1, cur; i <= q; ++i) {
read(cur);
printf("%d\n", ans[cur]);
}
return 0;
}