// #pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;
#define int long long
template<int SZ> struct Tree {
int n;
vector<int> adj[SZ];
void init(int N) {
n = N;
for (int i = 0; i < n-1; i++) {
int a, b; cin >> a >> b; a--, b--;
adj[a].push_back(b);
adj[b].push_back(a);
}
}
};
const int maxn = 10001;
const int maxk = 101;
int n;
int v[maxn];
int vis[maxk];
Tree<maxn> tree;
bitset<maxk> l[maxn][maxk], r[maxn][maxk];
void dfs(int nd, int par) {
if (vis[v[nd]]) return;
vis[v[nd]] = 1;
for (auto i : tree.adj[nd])
dfs(i, nd);
// set l
l[nd][v[nd]][v[nd]] = 1;
for (int i = v[nd]+1; i < maxk; i++)
if (l[nd][v[nd]][i-1])
for (auto j : tree.adj[nd])
if (j != par)
l[nd][v[nd]] |= l[j][i];
// set r
r[nd][v[nd]][v[nd]] = 1;
for (int i = v[nd]-1; i >= 0; i--)
if (r[nd][v[nd]][i+1])
for (auto j : tree.adj[nd])
if (j != par)
r[nd][v[nd]] |= r[j][i];
// set rest
for (int i = 0; i < maxk; i++)
if (i != v[nd] && l[nd][v[nd]][i])
r[nd][i] = r[nd][v[nd]];
for (int i = 0; i < maxk; i++)
if (i != v[nd] && r[nd][v[nd]][i])
l[nd][i] = l[nd][v[nd]];
vis[v[nd]] = 0;
}
// #define LOCAL
// #define CODEFORCES
signed main() {
#ifndef LOCAL
cin.tie(nullptr) -> ios::sync_with_stdio(false);
#endif
#ifdef LOCAL
freopen("main.in", "r", stdin);
#endif
cin >> n;
for (int i = 0; i < n; i++) {cin >> v[i]; v[i]--;}
tree.init(n);
dfs(0, 0);
int ans = 0;
for (int i = 0; i < maxk; i++)
ans += l[0][i].count();
cout << ans << endl;
return 0;
}
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
596 KB |
Output is correct |
2 |
Correct |
1 ms |
724 KB |
Output is correct |
3 |
Correct |
1 ms |
724 KB |
Output is correct |
4 |
Correct |
1 ms |
596 KB |
Output is correct |
5 |
Correct |
1 ms |
724 KB |
Output is correct |
6 |
Correct |
1 ms |
852 KB |
Output is correct |
7 |
Correct |
1 ms |
852 KB |
Output is correct |
8 |
Correct |
1 ms |
852 KB |
Output is correct |
9 |
Correct |
1 ms |
852 KB |
Output is correct |
10 |
Correct |
1 ms |
852 KB |
Output is correct |
11 |
Correct |
7 ms |
7892 KB |
Output is correct |
12 |
Correct |
6 ms |
3028 KB |
Output is correct |
13 |
Correct |
4 ms |
2132 KB |
Output is correct |
14 |
Correct |
3 ms |
1236 KB |
Output is correct |
15 |
Correct |
4 ms |
1236 KB |
Output is correct |
16 |
Correct |
3 ms |
1236 KB |
Output is correct |
17 |
Correct |
3 ms |
1620 KB |
Output is correct |
18 |
Correct |
3 ms |
1876 KB |
Output is correct |
19 |
Correct |
3 ms |
1364 KB |
Output is correct |
20 |
Correct |
3 ms |
1364 KB |
Output is correct |