#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;
#define int long long
template<typename T>
string tostr(const T& value) {
ostringstream oss;
oss << value;
return oss.str();
}
template<typename... Args>
string fstr(const string& format, Args... args) {
string result = format;
size_t pos = 0;
size_t argIndex = 0;
auto replaceArg = [&](const auto& arg) {
pos = result.find("{}", pos);
if (pos != string::npos) {
result.replace(pos, 2, tostr(arg));
++argIndex;
}
};
(replaceArg(args), ...);
return result;
}
/*
* Tree
*/
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);
}
}
};
/*
* LCA
* LCA-related functions
*/
template<int SZ, int LOG> struct LCA {
int n;
int bl[SZ][LOG];
int dpth[SZ];
void init_dfs(int nd, int par, vector<int>* adj) {
dpth[nd] = dpth[par]+1;
bl[nd][0] = par;
for (int i = 1; i < LOG; i++)
bl[nd][i] = bl[bl[nd][i-1]][i-1];
for (auto i : adj[nd])
if (i != par)
init_dfs(i, nd, adj);
}
void init(int N, Tree<SZ>& tree, int R = 0) {
n = N;
init_dfs(R, R, tree.adj);
}
int jmp(int nd, int k) {
for (int i = LOG-1; i >= 0; i--) if ((k >> i) & 1) nd = bl[nd][i];
return nd;
}
int lca(int a, int b) {
if (dpth[a] < dpth[b]) swap(a, b);
a = jmp(a, dpth[a] - dpth[b]);
if (a == b) return a;
for (int i = LOG-1; i >= 0; i--) {
int A = bl[a][i], B = bl[b][i];
if (A != B) a = A, b = B;
}
return bl[a][0];
}
bool onpath(int a, int b, int c) {
// whether c is on path from a --> b
int lc = lca(a, b);
return (dpth[lc] <= dpth[c]) && ((dpth[c] <= dpth[a] && jmp(a, dpth[a] - dpth[c]) == c) || (dpth[c] <= dpth[b] && jmp(b, dpth[b] - dpth[c]) == c));
}
};
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]];
// cout << "----------------- " << nd << endl;
// for (int i = 0; i < 4; i++) {
// cout << fstr("i: {}, l ", i);
// for (int j = 0; j < 4; j++)
// cout << l[nd][i][j];
// cout << endl;
// cout << fstr("i: {}, r ", i);
// for (int j = 0; j < 4; j++)
// cout << r[nd][i][j];
// cout << endl;
// }
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;
}
| # |
결과 |
실행 시간 |
메모리 |
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 |
10 ms |
7892 KB |
Output is correct |
| 12 |
Correct |
4 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 |
3 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 |
