#include <bits/stdc++.h>
using namespace std;
const int MAX_N = 1e5;
int state[MAX_N + 1], depth[MAX_N + 1], parent[MAX_N + 1];
vector<int> edges[MAX_N + 1], cities[MAX_N + 1];
set<int> stateEdges[MAX_N + 1];
struct DSU {
int comp;
int p[MAX_N + 1];
void init( int k ) {
for ( int u = 1; u <= k; u++ )
p[u] = u;
comp = k;
}
int find( int u ) {
if ( p[u] == u )
return u;
p[u] = find( p[u] );
return p[u];
}
void merge( int u, int v ) {
u = find( u );
v = find( v );
if ( u == v )
return;
comp--;
p[u] = v;
}
} dsu;
void dfs( int u, int p ) {
parent[u] = p;
depth[u] = depth[p] + 1;
for ( int v: edges[u] ) {
if ( v == p )
continue;
dfs( v, u );
}
}
int main() {
int n, k;
cin >> n >> k;
for ( int i = 0; i < n - 1; i++ ) {
int u, v;
cin >> u >> v;
edges[u].push_back( v );
edges[v].push_back( u );
}
for ( int u = 1; u <= n; u++ ) {
cin >> state[u];
cities[state[u]].push_back( u );
}
dfs( 1, 0 );
dsu.init( k );
for ( int s = 1; s <= k; s++ ) {
for ( int i = 0; i < cities[s].size() - 1; i++ ) {
int u = cities[s][i], v = cities[s][i + 1];
while ( depth[u] > depth[v] ) {
dsu.merge( state[u], state[parent[u]] );
u = parent[u];
}
while ( depth[v] > depth[u] ) {
dsu.merge( state[v], state[parent[v]] );
v = parent[v];
}
while ( u != v ) {
dsu.merge( state[u], state[parent[u]] );
u = parent[u];
dsu.merge( state[v], state[parent[v]] );
v = parent[v];
}
}
}
for ( int u = 1; u <= n; u++ )
state[u] = dsu.find( state[u] );
for ( int u = 1; u <= n; u++ ) {
for ( int v: edges[u] ) {
if ( state[u] == state[v] )
continue;
stateEdges[state[u]].insert( state[v] );
stateEdges[state[v]].insert( state[u] );
}
}
int leafs = 0;
for ( int s = 1; s <= k; s++ ) {
if ( stateEdges[s].size() == 1 )
leafs++;
}
if ( leafs != 0 )
return 1;
cout << (leafs == 0 ? 0 : leafs - 1);
return 0;
}
Compilation message
mergers.cpp: In function 'int main()':
mergers.cpp:67:28: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
67 | for ( int i = 0; i < cities[s].size() - 1; i++ ) {
| ~~^~~~~~~~~~~~~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Runtime error |
6 ms |
9708 KB |
Execution failed because the return code was nonzero |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Runtime error |
6 ms |
9708 KB |
Execution failed because the return code was nonzero |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Runtime error |
6 ms |
9708 KB |
Execution failed because the return code was nonzero |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
76 ms |
14836 KB |
Output is correct |
2 |
Runtime error |
177 ms |
27188 KB |
Execution failed because the return code was nonzero |
3 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Runtime error |
6 ms |
9708 KB |
Execution failed because the return code was nonzero |
2 |
Halted |
0 ms |
0 KB |
- |