#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pi;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<pi> vpi;
typedef vector<vi> vvi;
const int inf = 0x3f3f3f3f;
const ll linf = 123456789012345678;
const ll mod = 1000000007;
#pragma GCC optimize("Ofast")
#define all(x) x.begin(), x.end()
#define debug(x) cerr << #x << " = " << x << endl
#define sz(x) ((int)(x).size())
int n;
vvi adj;
vi k;
vvi adj2;
vi visited;
void dfs2(int node, bool back){
if(visited[node]==1) return;
visited[node] = 1;
if((!back && node < 60000) || (back && (node >= 60000 || node < 40000))){
for(int x : adj2[node]){
dfs2(x, back);
}
}
if(!back && node >= 60000) dfs2(node-20000+1, back);
if(back && node >= 40000 && node < 60000) dfs2(node+20000-1, back);
}
pair<vi, vi> dfs(int node, int par){
vector<pair<vi, vi>> results;
for(int x : adj[node]){
if(x == par) continue;
results.push_back(dfs(x, node));
}
for(int i = 0; i < sz(results); i++){
for(int x : results[i].first){
adj2[x+40000].push_back(i);
adj2[i+20000].push_back(x+40000);
}
for(int x : results[i].second){
adj2[x+60000].push_back(i+20000);
adj2[i].push_back(x+60000);
}
}
dfs2(60000 + k[node], false);
dfs2(40000 + k[node], true);
vi start, end;
for(int i = 0; i <= k[node]; i++){
if(visited[40000+i]) start.push_back(i);
}
for(int i = k[node]; i < 100; i++){
if(visited[60000+i]) end.push_back(i);
}
for(int i = 0; i < sz(results); i++){
for(int x : results[i].first){
adj2[x+40000].clear();
adj2[i+20000].clear();
visited[i] = 0;
visited[x+40000] = 0;
visited[x+60000-1] = 0;
}
for(int x : results[i].second){
adj2[x+60000].clear();
adj2[i].clear();
visited[i+20000] = 0;
visited[x+60000] = 0;
visited[x+40000+1] = 0;
}
}
visited[k[node]+60000] = 0;
visited[k[node]+40000+1] = 0;
visited[k[node]+40000] = 0;
visited[k[node]+60000-1] = 0;
return {start, end};
}
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
adj2 = vvi(100000);
visited = vi(100000);
cin >> n;
k = vi(n);
adj = vvi(n);
for(int i = 0; i < n; i++) {
cin >> k[i];
k[i]--;
}
for(int i = 0; i < n-1; i++){
int u, v;
cin >> u >> v;
u--, v--;
adj[u].push_back(v);
adj[v].push_back(u);
}
pair<vi, vi> p = dfs(0, -1);
cout << sz(p.first)*sz(p.second) << endl;
}
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
4 ms |
3064 KB |
Output is correct |
2 |
Correct |
4 ms |
3064 KB |
Output is correct |
3 |
Correct |
4 ms |
3064 KB |
Output is correct |
4 |
Correct |
4 ms |
3064 KB |
Output is correct |
5 |
Correct |
4 ms |
3064 KB |
Output is correct |
6 |
Correct |
4 ms |
3064 KB |
Output is correct |
7 |
Correct |
4 ms |
3192 KB |
Output is correct |
8 |
Correct |
4 ms |
3192 KB |
Output is correct |
9 |
Correct |
4 ms |
3064 KB |
Output is correct |
10 |
Correct |
4 ms |
3064 KB |
Output is correct |
11 |
Correct |
15 ms |
3704 KB |
Output is correct |
12 |
Correct |
15 ms |
3704 KB |
Output is correct |
13 |
Correct |
15 ms |
3676 KB |
Output is correct |
14 |
Correct |
36 ms |
6008 KB |
Output is correct |
15 |
Correct |
36 ms |
6136 KB |
Output is correct |
16 |
Correct |
36 ms |
6008 KB |
Output is correct |
17 |
Correct |
15 ms |
3704 KB |
Output is correct |
18 |
Correct |
15 ms |
3704 KB |
Output is correct |
19 |
Correct |
13 ms |
3832 KB |
Output is correct |
20 |
Correct |
13 ms |
3960 KB |
Output is correct |