#include <bits/stdc++.h>
#define N 1000000
using namespace std;
int range[N], n, k, city[N], ans[N], bef[N], pos, x[N];
vector<int> adj[N];
void dfs(int node, int par = 0){
pos++;
// cout << endl;
// cout << pos << '\n';
x[pos] = city[node];
if(bef[city[node]] == 0){
range[pos] = pos;
bef[city[node]] = pos;
ans[city[node]] = 1;
}
else{
range[pos] = 1e9;
int val = pos - 1;
while(val > bef[city[node]]){
// cout << val << ' ';
if(range[val] > bef[city[node]]) ans[city[node]] += ans[x[val]];
else{
ans[city[node]] += (ans[x[val]] - 1);
for(int i = val; i > bef[city[node]]; i--){
if(range[i] < bef[city[node]]) ans[x[i]] = ans[city[node]];
}
}
range[pos] = min(range[pos], range[val]);
val = range[val] - 1;
}
// if(val < bef[city[node]]){
//
// }
// cout << val << '\n';
range[pos] = min(range[pos], range[bef[city[node]]]);
}
// cout << ans[city[node]] << '\n';
for(int neigh : adj[node]){
if(neigh != par) dfs(neigh, node);
}
}
int main(){
ios_base::sync_with_stdio(0);
cin.tie(0);
cout.tie(0);
cin >> n >> k;
for(int i = 1; i < n; i++){
int a, b;
cin >> a >> b;
adj[a].push_back(b);
adj[b].push_back(a);
}
int leaf;
for(int i = 1; i <= n; i++){
cin >> city[i];
if(adj[i].size() == 1) leaf = i;
}
dfs(leaf);
int answer = 1e9;
for(int i = 1; i <= k; i++){
cout << ans[i] << ' ';
// answer = min(answer, ans[i] - 1);
}
// cout << '\n';
cout << answer;
}
Compilation message
capital_city.cpp: In function 'int main()':
capital_city.cpp:59:5: warning: 'leaf' may be used uninitialized in this function [-Wmaybe-uninitialized]
59 | dfs(leaf);
| ~~~^~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
12 ms |
23764 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
12 ms |
23764 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
127 ms |
46328 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
12 ms |
23764 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
