#include <iostream>
#include <bits/stdc++.h>
using namespace std;
typedef int64_t llo;
#define mp make_pair
#define pb push_back
#define a first
#define b second
vector<int> adj[500001];
int col[500001];
vector<int> cc[500001];
int par[500001][20];
int lev[500001];
int st[500001];
int endd[500001];
int co=0;
int dfs(int no,int par2=0,int lev2=0){
st[no]=co;
co+=1;
par[no][0]=par2;
lev[no]=lev2;
for(auto j:adj[no]){
if(j==par2){
continue;
}
dfs(j,no,lev2+1);
}
endd[no]=co-1;
}
int lca(int aa,int bb){
if(lev[aa]>lev[bb]){
swap(aa,bb);
}
int dif=lev[bb]-lev[aa];
// cout<<aa<<" "<<bb<<endl;
// cout<<dif<<endl;
for(int j=19;j>=0;j--){
if((1<<j)&dif){
// cout<<j<<","<<endl;
bb=par[bb][j];
}
}
// cout<<aa<<","<<bb<<endl;
if(aa==bb){
return aa;
}
for(int j=19;j>=0;j--){
if(par[aa][j]!=par[bb][j]){
aa=par[aa][j];
bb=par[bb][j];
}
}
return par[aa][0];
}
int par3[500001];
int find(int no){
if(par3[no]==no){
return no;
}
par3[no]=find(par3[no]);
return par3[no];
}
vector<int> proc[500001];
set<pair<int,int>> cur;
int dfs2(int no,int par=0){
if(cur.size()>0){
auto j=cur.lower_bound({st[no],0});
if(j==cur.end()){
}
else if((*j).a<=endd[no]){
par3[no]=par;
/// cout<<(*j).a<<" "<<(*j).b<<endl;
// cout<<no<<",,"<<endl;
}
}
for(auto j:proc[no]){
cur.insert({st[j],endd[j]});
}
auto tt=cur.find({st[no],endd[no]});
if(tt!=cur.end()){
cur.erase(tt);
}
for(auto j:adj[no]){
if(j==par){
continue;
}
dfs2(j,no);
}
}
int vis[500001];
vector<int> adj2[500001];
int ans=0;
int dfs3(int no,int par=-1){
vis[no]=1;
int co=0;
for(auto j:adj2[no]){
if(vis[j]==0){
dfs3(j,no);
co+=1;
}
}
if(par!=-1 and co==0){
ans+=1;
}
}
int main(){
ios_base::sync_with_stdio(false);
cin.tie(NULL);
memset(vis,0,sizeof(vis));
int n,k;
cin>>n>>k;
int aa,bb;
for(int i=0;i<n-1;i++){
cin>>aa>>bb;
adj[aa-1].pb(bb-1);
adj[bb-1].pb(aa-1);
}
for(int i=0;i<n;i++){
cin>>aa;
col[i]=aa;
cc[aa].pb(i);
}
dfs(0);
for(int i=0;i<n;i++){
par3[i]=i;
}
for(int i=0;i<n;i++){
for(int j=1;j<20;j++){
par[i][j]=par[par[i][j-1]][j-1];
}
}
for(int i=1;i<=k;i++){
aa=cc[i][0];
for(int j=1;j<cc[i].size();j++){
// cout<<i<<" "<<j<<" "<<aa<<" "<<cc[i][j]<<endl;
aa=lca(aa,cc[i][j]);
// cout<<i<<" "<<j<<" "<<aa<<" "<<cc[i][j]<<endl;
}
// cout<<aa<<","<<endl;
for(auto j:cc[i]){
// cout<<j<<":";
int x=j;
proc[aa].pb(x);
}
// cout<<endl;
// cout<<aa<<endl;
}
//cout<<44<<endl;
dfs2(0);
for(int i=0;i<n;i++){
find(i);
}
for(int i=0;i<n;i++){
for(auto j:adj[i]){
if(par3[j]!=par3[i]){
adj2[par3[j]].pb(par3[i]);
adj2[par3[i]].pb(par3[j]);
}
}
}
/* for(int i=0;i<n;i++){
cout<<par3[i]<<" ";
}
cout<<endl;*/
dfs3(0);
if(ans==0){
cout<<0<<endl;
}
else if(ans==1){
cout<<1<<endl;
}
else if(adj2[0].size()==1){
cout<<ans<<endl;
}
else{
cout<<ans-1<<endl;
}
return 0;
}
Compilation message
mergers.cpp: In function 'int dfs(int, int, int)':
mergers.cpp:29:1: warning: no return statement in function returning non-void [-Wreturn-type]
}
^
mergers.cpp: In function 'int dfs2(int, int)':
mergers.cpp:92:1: warning: no return statement in function returning non-void [-Wreturn-type]
}
^
mergers.cpp: In function 'int dfs3(int, int)':
mergers.cpp:108:1: warning: no return statement in function returning non-void [-Wreturn-type]
}
^
mergers.cpp: In function 'int main()':
mergers.cpp:138:16: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int j=1;j<cc[i].size();j++){
~^~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
32 ms |
49280 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
32 ms |
49280 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
32 ms |
49280 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
171 ms |
68224 KB |
Output is correct |
| 2 |
Incorrect |
184 ms |
72812 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
32 ms |
49280 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
