#include <bits/stdc++.h>
#define eb emplace_back
using namespace std;
const int MX = 2e5 + 5;
int N, K, C[MX];
int A, vis[MX], sz[MX];
vector<int> adj[MX], ele[MX];
int P[MX], cnt[MX], use[MX], clh[MX];
void d1(int n, int p) {
sz[n]=1; cnt[C[n]]++;
for (auto &i:adj[n])
if (i!=p&&!vis[i]) d1(i, n), sz[n]+=sz[i];
}
int d2(int n, int p, int c) {
for (auto &i:adj[n])
if (i!=p&&!vis[i]&&sz[i]>=c) return d2(i, n, c);
return n;
}
void d3(int n, int p) {
P[n]=p;
for (auto &i:adj[n])
if (i!=p&&!vis[i]) d3(i, n);
}
void d4(int n, int p) {
vector<int> st;
cnt[C[n]]=0;
for (auto &i:adj[n])
if (i!=p&&!vis[i]) d4(i, n);
}
void sol(int st) {
d1(st, 0);
int c=d2(st, 0, sz[st]/2+1);
d3(c, 0);
int fl=0, d=0;
vector<int> s; s.eb(c);
for (int i=0; i<s.size(); i++) {
int x=s[i];
if (use[x]) continue;
if (ele[C[x]].size()!=cnt[C[x]]) { fl=1; break; }
if (!clh[C[x]]) {
clh[C[x]]=1;
d++;
for (auto &j:ele[C[x]]) s.eb(j);
}
if (P[x]) s.eb(P[x]);
}
if (!fl) A=min(A, d);
for (auto &i:s) use[i]=0, clh[C[i]]=0;
d4(st);
s.clear();
vis[c]=1;
for (auto &i:adj[c]) if (!vis[i]) sol(i);
}
int main() {
cin.tie(0)->sync_with_stdio(0);
cin>>N>>K;
for (int i=1, u, v; i<N; i++) cin>>u>>v, adj[u].eb(v), adj[v].eb(u);
for (int i=1; i<=N; i++) cin>>C[i], ele[C[i]].eb(i);
A=MX;
sol(1);
cout<<A-1<<'\n';
return 0;
}
Compilation message
capital_city.cpp: In function 'void sol(int)':
capital_city.cpp:44:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
44 | for (int i=0; i<s.size(); i++) {
| ~^~~~~~~~~
capital_city.cpp:47:23: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
47 | if (ele[C[x]].size()!=cnt[C[x]]) { fl=1; break; }
| ~~~~~~~~~~~~~~~~^~~~~~~~~~~
capital_city.cpp:57:7: error: too few arguments to function 'void d4(int, int)'
57 | d4(st);
| ^
capital_city.cpp:30:6: note: declared here
30 | void d4(int n, int p) {
| ^~
