This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include<bits/stdc++.h>
using namespace std;
using ll = long long;
#define all(x) (x).begin(), (x).end()
const int N = 2e5 + 7;
int n, k;
vector<int> adj[N], col[N];
int c[N], cnt[N], pt[N];
int sz[N];
bool vis[N], rem[N];
int ans;
int dfs_sz(int v, int p = 0) {
sz[v] = 1;
for (int u : adj[v]) {
if (u != p && !rem[u])
sz[v] += dfs_sz(u, v);
}
return sz[v];
}
int get_centroid(int v, int tot, int p = 0) {
for (int u : adj[v]) {
if (u != p && !rem[u]) {
if (sz[u] * 2 > tot)
return get_centroid(u, tot, v);
}
}
return v;
}
void prep(int v, int p = 0) {
pt[v] = p;
cnt[c[v]]++;
for (int u : adj[v]) {
if (u != p && !rem[u])
prep(u, v);
}
}
void cls(int v, int p = 0) {
cnt[c[v]] = 0;
vis[c[v]] = 0;
for (int u : adj[v]) {
if (u != p && !rem[u])
cls(u, v);
}
}
void solve(int a = 1) {
int s = get_centroid(a, dfs_sz(a));
prep(s);
queue<int> q;
q.push(c[s]);
vis[c[s]] = 1;
bool ok = 1;
int used = 0;
while (q.size()) {
int x = q.front();
q.pop();
if (cnt[x] != col[x].size()) {
ok = 0;
break;
}
used++;
for (int v : col[x]) {
v = pt[v];
if (v > 0 && !vis[c[v]]) {
vis[c[v]] = 1;
q.push(c[v]);
}
}
}
if (ok)
ans = min(ans, used - 1);
cls(s);
rem[s] = 1;
for (int u : adj[s]) {
if (!rem[u])
solve(u);
}
}
int main() {
cin.tie(0)->sync_with_stdio(false);
cin >> n >> k;
for (int i = 1; i < n; i++) {
int u, v;
cin >> u >> v;
adj[u].push_back(v);
adj[v].push_back(u);
}
for (int i = 1; i <= n; i++) {
cin >> c[i];
col[c[i]].push_back(i);
}
ans = k - 1;
solve();
cout << ans << '\n';
return 0;
}
Compilation message (stderr)
capital_city.cpp: In function 'void solve(int)':
capital_city.cpp:69:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
69 | if (cnt[x] != col[x].size()) {
| ~~~~~~~^~~~~~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |