Submission #228766

#	Time	Username	Problem	Language	Result	Execution time	Memory
228766		kshitij_sodani	Mergers (JOI19_mergers)	C++17	100 / 100	1640 ms	221036 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=-1,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];

	if(dif>0){
		for(int j=19;j>=0;j--){
			if((1<<j)&dif){

				int cc=bb;
				bb=par[bb][j];


			}
		}
	}
	
	if(aa==bb){
		return aa;
	}
	
	for(int j=19;j>=0;j--){
		if(par[aa][j]!=par[bb][j] and par[aa][j]!=-1 and par[bb][j]!=-1){
			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 par2=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]=par2;
		}
	}
	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==par2){
			continue;
		}
		dfs2(j,no);
	}
}
set<int> adj2[500001];
int ans=0;

int main(){
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	int n,k;
	cin>>n>>k;
	int aa,bb;
	co=0;
	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);
/*	if(n==3 and k==2){
		while(true){
			continue;
		}
	}*/
	for(int i=0;i<n;i++){
		par3[i]=i;
	}
	for(int j=1;j<20;j++){
		for(int i=0;i<n;i++){
			if(par[i][j-1]==-1){
				par[i][j]=-1;
				continue;
			}
			if(par[par[i][j-1]][j-1]==-1){
				par[i][j]=-1;
				continue;
			}
			par[i][j]=par[par[i][j-1]][j-1];
			
		/*	if(lev[par[i][j]]-lev[par[i][j-1]]!=(1<<(j-1))){
				//cout<<i<<"                     "<<j<<endl;
			}*/
		}
	}
	for(int i=0;i<n;i++){
		//cout<<par[i][0]<<" ";
	}
	//cout<<endl;
	for(int i=1;i<=k;i++){
		if(cc[i].size()<=1){
			continue;
		}
		aa=lca(cc[i][0],cc[i][1]);
		for(int j=1;j<cc[i].size();j++){
			aa=lca(aa,cc[i][j]);
		}
		for(auto j:cc[i]){
			proc[aa].pb(j);
		}
//		//cout<<aa<<" ";
	}
//	//cout<<endl;
	dfs2(0);

	for(int i=0;i<n;i++){
		find(i);
	}
	for(int i=0;i<n;i++){
		//cout<<par3[i]<<" ";
	}
	//cout<<endl;
	for(int i=0;i<n;i++){
		for(auto j:adj[i]){
			if(find(par3[j])!=find(par3[i])){
				adj2[find(par3[j])].insert(find(par3[i]));
				adj2[find(par3[i])].insert(find(par3[j]));
			}
		}
	}

	for(int i=0;i<n;i++){
		if(adj2[i].size()==1){
			ans+=1;
		}
	}
	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 lca(int, int)':
mergers.cpp:40:9: warning: unused variable 'cc' [-Wunused-variable]
     int cc=bb;
         ^~
mergers.cpp: In function 'int dfs2(int, int)':
mergers.cpp:94:1: warning: no return statement in function returning non-void [-Wreturn-type]
 }
 ^
mergers.cpp: In function 'int main()':
mergers.cpp:152:16: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int j=1;j<cc[i].size();j++){
               ~^~~~~~~~~~~~~

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