/* in the name of allah */
#include<bits/stdc++.h>
using namespace std;
#define endl '\n'
#define pb push_back
#define F first
#define S second
#define mk make_pair
typedef long long ll;
const int N=5e5+7,lg=21;
int n,k,st[N],fn[N],t,s[N],mark[N],sum[N],ans,h[N],mn[N],sp[N][lg],P[N];
vector<int>vec[N],ver[N],x;
int lca(){
int m=N,g;
for(auto u:x){
if(h[u]<m){
m=h[u];
g=u;
}
}
for(int i=lg-1;i>=0;i--){
for(int j=0;j<x.size();j++){
int u=x[j];
if(h[sp[u][i]]>=m){
x[j]=sp[u][i];
}
}
}
int f=1;
for(auto u:x){
if(u!=g)f=0;
}
if(f)return g;
for(int i=lg-1;i>=0;i--){
f=1;
for(int j=1;j<x.size();j++){
if(sp[x[j]][i]==sp[x[0]][i])f=0;
}
if(f){
for(int j=0;j<x.size();j++){
x[j]=sp[x[j]][i];
}
}
}
return sp[x[0]][0];
}
void dfs1(int v=1,int p=0){
h[v]=h[p]+1;
sp[v][0]=p;
st[v]=t++;
for(auto u:vec[v]){
if(u!=p) dfs1(u,v);
}
fn[v]=t++;
}
void dfs2(int v=1,int p=0){
mn[v]=P[s[v]];
for(auto u:vec[v]){
if(u==p)continue;
dfs2(u,v);
sum[v]+=sum[u];
mn[v]=min(mn[v],mn[u]);
}
if(v!=1 && h[v]<=mn[v])mark[v]=1;
}
void dfs3(int v=1,int p=0){
if(mark[v] && (sum[v]==1 || sum[v]==sum[1]))ans++;
for(auto u:vec[v]){
if(u==p)continue;
dfs3(u,v);
}
}
int main(){
ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
cin>>n>>k;
int u,v;
for(int i=0;i<n-1;i++){
cin>>u>>v;
vec[u].pb(v);
vec[v].pb(u);
}
for(int i=1;i<=n;i++){
cin>>s[i];
ver[s[i]].pb(i);
}
dfs1();
for(int j=1;j<lg;j++){
for(int i=1;i<=n;i++){
sp[i][j]=sp[sp[i][j-1]][j-1];
}
}
for(int i=1;i<=k;i++){
x.clear();
for(auto u:ver[i])x.pb(u);
P[i]=lca();
}
dfs2();
dfs3();
cout<<ans/2+(ans%2==1);
}
Compilation message
mergers.cpp: In function 'int lca()':
mergers.cpp:25:16: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
25 | for(int j=0;j<x.size();j++){
| ~^~~~~~~~~
mergers.cpp:39:16: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
39 | for(int j=1;j<x.size();j++){
| ~^~~~~~~~~
mergers.cpp:43:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
43 | for(int j=0;j<x.size();j++){
| ~^~~~~~~~~
mergers.cpp:17:10: warning: 'g' may be used uninitialized in this function [-Wmaybe-uninitialized]
17 | int m=N,g;
| ^
# |
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 |
12 ms |
23764 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Incorrect |
76 ms |
39936 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 |
- |