树上数颜色

题目链接

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
#include <bits/stdc++.h>
/*
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/hash_policy.hpp>
*/
using namespace std;

const double eps = 1e-10;
const double pi = 3.1415926535897932384626433832795;
const double eln = 2.718281828459045235360287471352;

#define f(i, a, b) for (int i = a; i <= b; i++)
#define scan(x) scanf("%d", &x)
#define mp make_pair
#define pb push_back
#define lowbit(x) (x & (-x))

#define fi first
#define se second
#define SZ(x) int((x).size())
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
#define summ(a) (accumulate(all(a), 0ll))

typedef unsigned long long ull;
typedef pair<int, int> pii;
typedef vector<int> vi;

using ll = long long;

ll n, m, ans[(int)1e5 + 9], c[(int)1e5 + 9], cnt[(int)1e5 + 9], sz[(int)1e5 + 9], son[(int)1e5 + 9], res;
vector<int> g[(int)1e5 + 9];

void dfs(int u, int fa)
{
for (auto it : g[u])
{
if (it != fa)
{
dfs(it, u), sz[u] += sz[it], son[u] = sz[son[u]] < sz[it] ? it : son[u];
}
}
sz[u] += 1;
}

void cal(int u, int fa, int val, int bson)
{
cnt[c[u]] += val;
if (val == 1)
res += cnt[c[u]] == 1;
else
res -= cnt[c[u]] == 0;
for (auto it : g[u])
if (it != fa && it != bson)
cal(it, u, val, bson);
}

void dfs2(int u, int fa, int opt)
{
for (auto it : g[u])
{
if (it != son[u] && it != fa)
{
dfs2(it, u, 1);
}
}
if (son[u])
dfs2(son[u], u, 0);
cal(u, fa, 1, son[u]);
ans[u] = res;
if (opt)
cal(u, fa, -1, 0);
}

int main()
{
ios::sync_with_stdio(false);
cin.tie(0);
cin >> n;
for (int i = 1, x, y; i < n; ++i)
{
cin >> x >> y;
g[x].push_back(y), g[y].push_back(x);
}
f(i, 1, n) cin >> c[i];
dfs(1, 0);
dfs2(1, 0, 0);
cin >> m;
for (int i = 1, x; i <= m; ++i)
{
cin >> x;
cout << ans[x] << "\n";
}
return 0;
}