Submission #228720

#	Time	Username	Problem	Language	Result	Execution time	Memory
228720		kshitij_sodani	Mergers (JOI19_mergers)	C++17	0 / 100	266 ms	89968 KiB

mergers

#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).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);
	}
}
set<int> adj2[500001];
int ans=0;
//int cop=0;
//int stt=1;
int dfs3(int no,int par3=0){
	int co=0;
	int coo=adj2[no].size()-1;
	if(no==0){
		coo+=1;
	}
/*	if(no!=0 and st==1){
		cop+=1;
	}*/
	/*if(coo>1){
		stt=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;
	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]].insert(par3[i]);
				adj2[par3[i]].insert(par3[j]);
			}
		}
	}
	/*for(int i=0;i<n;i++){
		cout<<par3[i]<<" ";
	}
	cout<<endl;*/
	dfs3(0);
cout<<(ans+1)/2<<endl;

	return 0;
}

Compilation message (stderr)

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:90:1: warning: no return statement in function returning non-void [-Wreturn-type]
 }
 ^
mergers.cpp: In function 'int dfs3(int, int)':
mergers.cpp:96:6: warning: unused variable 'co' [-Wunused-variable]
  int co=0;
      ^~
mergers.cpp:115:1: warning: no return statement in function returning non-void [-Wreturn-type]
 }
 ^
mergers.cpp: In function 'int main()':
mergers.cpp:144:16: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int j=1;j<cc[i].size();j++){
               ~^~~~~~~~~~~~~

• Subtask 1: 0 / 10
• Subtask 2: 0 / 24
• Subtask 3: 0 / 14
• Subtask 4: 0 / 22
• Subtask 5: 0 / 30