#include <bits/stdc++.h>
using namespace std;
vector<vector<int>> adj;
int C[200005];
bool vis[200005];
int sz[200005];
//기억을 되살려서 센트 짜는중 sibal
vector<vector<int>> B;
int get_sz(int c, int p) {
sz[c] = 1;
for(int n : adj[c]) {
if(n==p || vis[n]) continue;
sz[c] += get_sz(n, c);
}
return sz[c];
}
int get_cent(int c, int p, int s) {
for(int n : adj[c]) {
if(n==p || vis[n]) continue;
if(sz[n] > s/2) return get_cent(n, c, s);
}
return c;
}
int par[200005];
int id[200005];
bool vis2[200005];
int cid[200005];
void dfs(int c, int p, int d) {
int cnt = 0;
par[c] = p;
if(d != -1) id[c] = d;
vis2[c] = false;
for(int n : adj[c]) {
if(n==p || vis[n]) continue;
if(d != -1) dfs(n, c, d);
else dfs(n, c, cnt);
cnt++;
}
}
set<int> Col;
void dfs2(int c, int p, int d) {
if(cid[C[c]]==d) Col.insert(C[c]);
for(int n : adj[c]) {
if(n==p || vis[n]) continue;
dfs2(n, c, d);
}
}
int DnC(int c) {
if(Col.size()==0) return 1e9;
assert(!vis[c]);
get_sz(c, -1);
int cen = get_cent(c, -1, sz[c]);
dfs(cen, -1, -1);
set<int> S;
int val = 1e9;
vis[cen] = true;
if(Col.find(C[cen])!=Col.end()) {
set<int> S;
queue<int> Q;
bool isPos = true;
S.insert(C[cen]);
for(int n : B[C[cen]]) Q.push(n);
while(!Q.empty()) {
int c = Q.front();
Q.pop();
if(vis2[c]) continue;
while(c != -1 && !vis2[c]) {
if(Col.find(C[c])==Col.end()) {
isPos = false;
break;
}
if(S.find(C[c])==S.end()) {
S.insert(C[c]);
for(int n : B[C[c]]) Q.push(n);
}
vis2[c] = true;
c = par[c];
}
if(!isPos) break;
}
if(isPos) val = min(val, (int)S.size()-1);
}
vector<set<int>> S2;
int cnt = 0;
for(int n : adj[cen]) {
if(!vis[n]) cnt++;
}
for(int col : Col) {
if(col == C[cen]) continue;
int d = -1;
for(int n : B[col]) {
if(d==-1) d = id[n];
else if(d >= 0 && d != id[n]) d = -2;
}
cid[col] = d;
}
cnt = 0;
for(int n : adj[cen]) {
if(!vis[n]) {
Col = set<int> {};
dfs2(n, -1, cnt);
val = min(val, DnC(n));
cnt++;
}
}
return val;
}
signed main() {
cin.sync_with_stdio(false);
cin.tie(0);
cout.tie(0);
int N, K;
cin >> N >> K;
int i, j;
adj.resize(N);
for(i=0;i<N-1;i++) {
int a, b;
cin >> a >> b;
adj[a-1].push_back(b-1);
adj[b-1].push_back(a-1);
}
B.resize(K);
for(i=0;i<N;i++) {
cin >> C[i];
C[i]--;
B[C[i]].push_back(i);
}
for(i=0;i<K;i++) {
if(B[i].size()==1) {
cout << "0";
return 0;
}
}
for(i=0;i<K;i++) Col.insert(i);
cout << DnC(0);
}
Compilation message
capital_city.cpp: In function 'int main()':
capital_city.cpp:115:12: warning: unused variable 'j' [-Wunused-variable]
115 | int i, j;
| ^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
0 ms |
340 KB |
Output is correct |
| 2 |
Incorrect |
0 ms |
340 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
0 ms |
340 KB |
Output is correct |
| 2 |
Incorrect |
0 ms |
340 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
207 ms |
37308 KB |
Output is correct |
| 2 |
Correct |
111 ms |
37668 KB |
Output is correct |
| 3 |
Correct |
189 ms |
37016 KB |
Output is correct |
| 4 |
Correct |
107 ms |
37672 KB |
Output is correct |
| 5 |
Correct |
235 ms |
34544 KB |
Output is correct |
| 6 |
Incorrect |
112 ms |
38236 KB |
Output isn't correct |
| 7 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
0 ms |
340 KB |
Output is correct |
| 2 |
Incorrect |
0 ms |
340 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
