#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
using namespace std;
#define int long long
#define mp make_pair
#define pb push_back
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("avx2")
constexpr int INF = (int)2e18;
constexpr int N = (int)2e5 + 111;
constexpr int md = (int)998244353;
constexpr int mdPower = (int)1e9+7 - 1;
constexpr double eps = 1e-7;
typedef __int128 _uint;
typedef long long ll;
mt19937_64 rnd(time(nullptr));
void add(int& a,int b){
a += b; if(a >= md) a -= md;
}
void sub(int& a,int b){
a -= b; while(a < 0) a += md;
}
void add(__int128& a,int b){
a += b;
}
void sub(__int128& a,int b){
a -= b;
}
int bpow(int a,int b){
if(b == 0)
return 1;
if(b % 2 == 0){
int t = bpow(a,b>>1);
return 1ll*t*t%md;
}
return 1ll * a*bpow(a,b-1) % md;
}
int inv(int a){
return bpow(a,md-2);
}
int fac[N],invfac[N];
void init(){
fac[0] = 1;
for(int i = 1; i < N; i++){
fac[i] = (fac[i-1] * i) % md;
}
invfac[N-1] = inv(fac[N-1]);
for(int i = N-2; i >= 0; i--){
invfac[i] = (invfac[i+1] * (i+1))%md;
}
return;
}
//int C(int n,int k){
// if(k > n)
// return 0;
// return fac[n] * invfac[k] % md * invfac[n-k] % md;
//}
//
//int A(int n,int k){
// if(k > n)
// return 0;
// return fac[n] * invfac[n-k] % md;
//}
vector<int> g[N];
int col[N];
vector<int> cities[N];
int sz[N];
int ans = INF;
bool used[N];
bool choosen[N];
int p[N];
int L[N];
int d[N];
void dfs1(int v,int pr = -1){
sz[v] = 1;
for(auto& to : g[v]){
if(to == pr || used[to])
continue;
dfs1(to,v);
sz[v] += sz[to];
}
return;
}
int getCentroid(int v,int n,int pr = -1){
for(auto& to : g[v]){
if(to == pr || used[to])
continue;
if(sz[to] > n/2)
return getCentroid(to,n,v);
}
return v;
}
int CNT = 0;
void dfs3(int v,int pr = -1){
L[v] = CNT;
for(auto& to : g[v]){
if(to == pr || used[to])
continue;
p[to] = v;
d[to] = d[v] + 1;
dfs3(to,v);
}
return;
}
void solve(int v,int pr = -1){
dfs1(v,pr);
int C = getCentroid(v,sz[v],pr);
// cerr << "C = " << C <<s "\n";
used[C] = true;
int cur_col = col[C];
dfs1(C);
CNT++;
d[C] = 0;
dfs3(C);
// cerr << "ok\n";
if(!choosen[cur_col]){
set<int> added_colors;
added_colors.insert(cur_col);
set<pair<int,int>> q;
bool ok = true;
for(auto& x : cities[cur_col]){
q.insert(mp(-d[x],x));
if(L[x] != CNT){
ok = false;
break;
}
}
while(!q.empty() && ok){
int v = q.begin()->second;
q.erase(q.begin());
// cerr << "v = " << v << "\n";
if(choosen[col[v]]){
ok = false;
break;
}
if(!added_colors.count(col[v])){
for(auto& x : cities[col[v]]){
if(x != v)
q.insert(mp(-d[x],x));
if(L[x] != CNT){
ok = false;
}
}
added_colors.insert(col[v]);
}
if(p[v] > 0)
q.insert(mp(-d[p[v]],p[v]));
}
choosen[cur_col] = 1;
if(ok)
ans = min(ans,(int)added_colors.size()-1);
// cerr << "v = " << v << ", centroid = " << C << "\n";
}
f:;
for(auto& to : g[C]){
if(!used[to]){
solve(to,v);
}
}
return;
}
void solve(){
int n,k;
cin >> n >> k;
for(int i = 1; i < n; i++){
int a,b;
cin >> a >> b;
g[a].pb(b);
g[b].pb(a);
}
for(int i = 1; i <= n; i++){
cin >> col[i];
cities[col[i]].pb(i);
}
solve(1);
cout << ans << "\n";
return;
}
signed main(){
ios::sync_with_stdio(false);cin.tie(nullptr);cout.tie(nullptr);
init();
int tests = 1;
// cin >> tests;
for(int test = 1; test <= tests; test++){
// cerr << "test = " << test << "\n";
solve();
}
return 0;
}
/**
**/
Compilation message
capital_city.cpp: In function 'void solve(long long int, long long int)':
capital_city.cpp:173:5: warning: label 'f' defined but not used [-Wunused-label]
173 | f:;
| ^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
8 ms |
12884 KB |
Output is correct |
| 2 |
Correct |
9 ms |
12884 KB |
Output is correct |
| 3 |
Correct |
18 ms |
12884 KB |
Output is correct |
| 4 |
Correct |
9 ms |
12812 KB |
Output is correct |
| 5 |
Correct |
8 ms |
12884 KB |
Output is correct |
| 6 |
Correct |
14 ms |
12884 KB |
Output is correct |
| 7 |
Execution timed out |
3063 ms |
12884 KB |
Time limit exceeded |
| 8 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
8 ms |
12884 KB |
Output is correct |
| 2 |
Correct |
9 ms |
12884 KB |
Output is correct |
| 3 |
Correct |
18 ms |
12884 KB |
Output is correct |
| 4 |
Correct |
9 ms |
12812 KB |
Output is correct |
| 5 |
Correct |
8 ms |
12884 KB |
Output is correct |
| 6 |
Correct |
14 ms |
12884 KB |
Output is correct |
| 7 |
Execution timed out |
3063 ms |
12884 KB |
Time limit exceeded |
| 8 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
629 ms |
36080 KB |
Output is correct |
| 2 |
Correct |
233 ms |
36252 KB |
Output is correct |
| 3 |
Correct |
575 ms |
35976 KB |
Output is correct |
| 4 |
Correct |
237 ms |
36192 KB |
Output is correct |
| 5 |
Execution timed out |
3070 ms |
35568 KB |
Time limit exceeded |
| 6 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
8 ms |
12884 KB |
Output is correct |
| 2 |
Correct |
9 ms |
12884 KB |
Output is correct |
| 3 |
Correct |
18 ms |
12884 KB |
Output is correct |
| 4 |
Correct |
9 ms |
12812 KB |
Output is correct |
| 5 |
Correct |
8 ms |
12884 KB |
Output is correct |
| 6 |
Correct |
14 ms |
12884 KB |
Output is correct |
| 7 |
Execution timed out |
3063 ms |
12884 KB |
Time limit exceeded |
| 8 |
Halted |
0 ms |
0 KB |
- |
