Submission #235279

#	Time	Username	Problem	Language	Result	Execution time	Memory
235279		kshitij_sodani	Capital City (JOI20_capital_city)	C++17	100 / 100	1186 ms	68832 KiB

capital_city

#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
int n,k;
int aa,bb;
int it[200001];
set<int> adj[200001];
int si[200001];

vector<int> col[200001];
vector<int> col2[200001];

vector<int> rem;
vector<int> rem2;


int vis[200001];
int vis2[200001];
int par[200001];


int count2=0;
int cur=0;
int dfs(int no,int par=-1){
	si[no]=1;
	rem.pb(it[no]-1);
	rem2.pb(no);
	col[it[no]-1].pb(no);
	for(auto j:adj[no]){
		if(j==par){
			continue;
		}
		dfs(j,no);
		si[no]+=si[j];
	}
}
int find(int no,int par=-1){
	for(auto j:adj[no]){
		if(j==par){
			continue;
		}
		if(si[j]>si[cur]/2){
			return find(j,no);
		}
	}
	return no;
}
int dfs2(int no,int par2=-1){
	par[no]=par2;
	for(auto j:adj[no]){
		if(j==par2){
			continue;
		}
		dfs2(j,no);
	}
}
int ma;
int dec(int no){
	cur=no;
	dfs(cur);
	int cen=find(no);
	count2=0;
	dfs2(cen);
	queue<int> ss;
	ss.push(cen);
	int st=1;
	int ans=0;
	vis2[cen]=1;
	while(ss.size()){
		int x=ss.front();
		ss.pop();
		if(vis[it[x]-1]==0){
			vis[it[x]-1]=1;
			ans+=1;
			if(col[it[x]-1].size()!=col2[it[x]-1].size()){
				st=0;
				break;
			}
			for(auto j:col[it[x]-1]){
				if(j!=x){
					ss.push(j);
				}
			}
		}
		x=par[x];
		while(x!=-1){
			if(vis2[x]==1){
				break;
			}
			ss.push(x);
			vis2[x]=1;
			x=par[x];
		}
	}
	if(st==1){
		ma=min(ma,ans-1);
	}
	for(auto j:rem){
		col[j].clear();
		vis[j]=0;
	}
	for(auto j:rem2){
		vis2[j]=0;
	}
	rem2.clear();
	rem.clear();
	for(auto j:adj[cen]){
		adj[j].erase(cen);
	}
	for(auto j:adj[cen]){
		dec(j);
	}
}
int main(){
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cin>>n>>k;
	for(int i=0;i<n-1;i++){
		cin>>aa>>bb;
		adj[aa-1].insert(bb-1);
		adj[bb-1].insert(aa-1);
	}
	for(int i=0;i<n;i++){
		cin>>it[i];
		col2[it[i]-1].pb(i);
	}
	ma=n-1;
	dec(0);
	cout<<ma<<endl;

	return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'int dfs(int, int)':
capital_city.cpp:40:1: warning: no return statement in function returning non-void [-Wreturn-type]
 }
 ^
capital_city.cpp: In function 'int dfs2(int, int)':
capital_city.cpp:60:1: warning: no return statement in function returning non-void [-Wreturn-type]
 }
 ^
capital_city.cpp: In function 'int dec(int)':
capital_city.cpp:117:1: warning: no return statement in function returning non-void [-Wreturn-type]
 }
 ^

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59