#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];
for(int j=19;j>=0;j--){
if((1<<j)&dif){
// cout<<j<<","<<endl;
bb=par[bb][j];
}
}
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).b<=endd[no]){
par3[no]=par;
}
}
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);
}
}
set<int> adj2[500001];
int ans=0;
int dfs3(int no,int par3=0){
for(auto j:adj2[no]){
if(j!=par3){
dfs3(j,no);
}
}
if(adj2[no].size()==1){
ans+=1;
}
}
int main(){
ios_base::sync_with_stdio(false);
cin.tie(NULL);
int n,k;
cin>>n>>k;
int aa,bb;
cur.clear();
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++){
aa=lca(aa,cc[i][j]);
}
for(auto j:cc[i]){
proc[aa].pb(j);
}
}
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]].insert(par3[i]);
adj2[par3[i]].insert(par3[j]);
}
}
}
dfs3(0);
cout<<(ans+1)/2<<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:84:1: warning: no return statement in function returning non-void [-Wreturn-type]
}
^
mergers.cpp: In function 'int dfs3(int, int)':
mergers.cpp:96:1: warning: no return statement in function returning non-void [-Wreturn-type]
}
^
mergers.cpp: In function 'int main()':
mergers.cpp:126: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 |
Correct |
39 ms |
59128 KB |
Output is correct |
| 2 |
Incorrect |
44 ms |
59256 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
39 ms |
59128 KB |
Output is correct |
| 2 |
Incorrect |
44 ms |
59256 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
39 ms |
59128 KB |
Output is correct |
| 2 |
Incorrect |
44 ms |
59256 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
175 ms |
77932 KB |
Output is correct |
| 2 |
Correct |
263 ms |
88176 KB |
Output is correct |
| 3 |
Incorrect |
43 ms |
59768 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
39 ms |
59128 KB |
Output is correct |
| 2 |
Incorrect |
44 ms |
59256 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
