//#pragma GCC optimize("O3,unroll-loops")
//#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include <bits/stdc++.h>
//#define int long long
#define ll long long
#define fi first
#define se second
#define pb push_back
#define all(lmao) lmao.begin(), lmao.end()
using namespace std;
typedef pair<int,int> pii;
typedef tuple<int,int,int> tp;
const int N = 5e5 + 5;
const ll oo = 1e16;
const int mod = 1e9 + 7;
const int base = 23;
int n, k, X[N], Y[N], col[N], a[N], f[N], dp[N], in[N], en[N], demin, b[N];
vector<int> p[N], g[N], vr[N], m;
bool fl[N];
void pre(int u,int v){
f[u] = a[u];
dp[u] = 1;
in[u] = ++demin;
b[demin] = u;
for(auto j : p[u]){
if(j == v) continue;
pre(j, u);
f[u] += f[j];
dp[u] += dp[j];
}
en[u] = demin;
}
void dfs(int u,int v){
if(dp[u] == f[u]){
m.pb(u);
return;
}
for(auto j : p[u]){
if(j == v) continue;
dfs(j, u);
}
}
signed main(){
ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#define task "v"
if(fopen(task ".inp","r")){
freopen(task ".inp","r",stdin);
freopen(task ".out","w",stdout);
}
cin >> n >> k;
for(int i = 1; i <= n - 1; i ++){
cin >> X[i] >> Y[i];
p[X[i]].pb(Y[i]);
p[Y[i]].pb(X[i]);
}
int root = 1;
for(int i = 1; i <= n; i ++){
if((int)p[i].size() >= 2) root = i;
int x; cin >> x;
vr[x].pb(i);
col[i] = x;
}
for(int i = 1; i < n; i ++){
if(col[X[i]] == col[Y[i]]){
g[X[i]].pb(Y[i]);
g[Y[i]].pb(X[i]);
}
}
if(n == 1){
cout << 0;
return 0;
}
if(n == 2){
if(col[1] == col[2]){
cout << 0 << "\n";
}else cout << 1 << "\n";
return 0;
}
for(int i = 1; i <= k; i ++){
queue<int> q;
q.push(vr[i][0]);
fl[vr[i][0]] = true;
while(!q.empty()){
int u = q.front();
q.pop();
for(auto j : g[u]){
if(fl[j]) continue;
fl[j] = true;
q.push(j);
}
}
bool ff = true;
for(auto j : vr[i]){
if(!fl[j]){
ff = false;
break;
}
}
if(ff) for(auto j : vr[i]) a[j] = 1;
}
pre(root, 0);
dfs(root, 0);
int ans = 0;
for(auto j : m){
set<int> st;
for(int i = in[j]; i <= en[j]; i ++) st.insert(col[b[i]]);
ans += ((int)st.size() - 1);
}
ans += max(0, (int)m.size() - 1);
cout << ans;
}
Compilation message
mergers.cpp: In function 'int main()':
mergers.cpp:55:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
55 | freopen(task ".inp","r",stdin);
| ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~
mergers.cpp:56:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
56 | freopen(task ".out","w",stdout);
| ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Incorrect |
21 ms |
49500 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Incorrect |
21 ms |
49500 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Incorrect |
21 ms |
49500 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Incorrect |
40 ms |
54360 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Incorrect |
21 ms |
49500 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |