#include <bits/stdc++.h>
using namespace std;
#ifdef MIKU
string dbmc = "\033[1;38;2;57;197;187m", dbrs = "\033[0m";
#define debug(x...) cout << dbmc << "[" << #x << "]: ", dout(x)
void dout() { cout << dbrs << endl; }
template <typename T, typename ...U>
void dout(T t, U ...u) { cout << t << (sizeof...(u) ? ", " : ""); dout(u...); }
#else
#define debug(...) 39
#endif
#define fs first
#define sc second
#define mp make_pair
#define FOR(i, j, k) for (int i = j, Z = k; i < Z; i++)
using ll = long long;
typedef pair<int, int> pii;
const int MXN = 10305;
int n, a[MXN];
vector<int> edge[MXN];
vector<int> vl[MXN], vr[MXN];
namespace GRAPH {
vector<int> edge[MXN];
vector<int> lk;
vector<int> al, ar;
bitset<MXN> b;
int nc = 105;
void MAKE_L(int sr) {
for (auto &id : lk) {
for (auto &i : vl[id]) edge[nc].push_back(i);
for (auto &i : vr[id]) edge[i].push_back(nc);
nc++;
}
al.clear();
b.reset();
al.push_back(sr);
b[sr] = 1;
for (int j = 0; j < al.size(); j++) {
int id = al[j];
for (auto &i : edge[id]) {
if (b[i]) continue;
b[i] = true;
al.push_back(i);
}
}
FOR(i, 0, nc) edge[i].clear();
nc = 105;
}
void MAKE_R(int sr) {
for (auto &id : lk) {
for (auto &i : vl[id]) edge[i].push_back(nc);
for (auto &i : vr[id]) edge[nc].push_back(i);
nc++;
}
ar.clear();
b.reset();
ar.push_back(sr);
b[sr] = 1;
for (int j = 0; j < ar.size(); j++) {
int id = ar[j];
for (auto &i : edge[id]) {
if (b[i]) continue;
b[i] = true;
ar.push_back(i);
}
}
FOR(i, 0, nc) edge[i].clear();
nc = 105;
}
void DO(int sr) {
MAKE_L(sr);
MAKE_R(sr + 1);
}
}
void MV(vector<int> &v, vector<int> &a) {
for (auto &i : a) {
if (i < 105) v.push_back(i);
}
a.clear();
}
void DFS(int id, int par) {
for (auto &i : edge[id]) {
if (i == par) continue;
DFS(i, id);
}
GRAPH::lk.clear();
for (auto &i : edge[id]) {
if (i == par) continue;
GRAPH::lk.push_back(i);
}
GRAPH::DO(a[id]);
MV(vl[id], GRAPH::al);
MV(vr[id], GRAPH::ar);
// debug(id);
// for (auto &i : vl[id]) cout << i << ' ';
// cout << endl;
// for (auto &i : vr[id]) cout << i << ' ';
// cout << endl;
}
void miku() {
int x, y;
cin >> n;
FOR(i, 1, n + 1) cin >> a[i];
FOR(i, 1, n) {
cin >> x >> y;
edge[x].push_back(y);
edge[y].push_back(x);
}
DFS(1, 0);
cout << vl[1].size() * vr[1].size() << '\n';
}
int32_t main() {
cin.tie(0) -> sync_with_stdio(false);
cin.exceptions(cin.failbit);
miku();
return 0;
}
Compilation message
uzastopni.cpp: In function 'void GRAPH::MAKE_L(int)':
uzastopni.cpp:43:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
43 | for (int j = 0; j < al.size(); j++) {
| ~~^~~~~~~~~~~
uzastopni.cpp: In function 'void GRAPH::MAKE_R(int)':
uzastopni.cpp:65:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
65 | for (int j = 0; j < ar.size(); j++) {
| ~~^~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
1624 KB |
Output is correct |
| 2 |
Correct |
1 ms |
1372 KB |
Output is correct |
| 3 |
Correct |
2 ms |
1372 KB |
Output is correct |
| 4 |
Correct |
1 ms |
1432 KB |
Output is correct |
| 5 |
Correct |
1 ms |
1332 KB |
Output is correct |
| 6 |
Correct |
1 ms |
1372 KB |
Output is correct |
| 7 |
Correct |
1 ms |
1368 KB |
Output is correct |
| 8 |
Correct |
1 ms |
1432 KB |
Output is correct |
| 9 |
Correct |
1 ms |
1372 KB |
Output is correct |
| 10 |
Correct |
1 ms |
1368 KB |
Output is correct |
| 11 |
Correct |
7 ms |
2396 KB |
Output is correct |
| 12 |
Correct |
10 ms |
2396 KB |
Output is correct |
| 13 |
Correct |
7 ms |
2396 KB |
Output is correct |
| 14 |
Correct |
13 ms |
5724 KB |
Output is correct |
| 15 |
Correct |
15 ms |
5720 KB |
Output is correct |
| 16 |
Correct |
20 ms |
5724 KB |
Output is correct |
| 17 |
Correct |
10 ms |
2300 KB |
Output is correct |
| 18 |
Correct |
7 ms |
2392 KB |
Output is correct |
| 19 |
Correct |
6 ms |
2392 KB |
Output is correct |
| 20 |
Correct |
6 ms |
2396 KB |
Output is correct |
