Submission #924699

#TimeUsernameProblemLanguageResultExecution timeMemory
924699NintsiChkhaidzeCapital City (JOI20_capital_city)C++17
100 / 100
440 ms43192 KiB
#include <bits/stdc++.h>
#define ll long long
#define s second
#define f first
#define pb push_back
#define pii pair <int,int>
#define left (h<<1),l,(l + r)/2
#define right ((h<<1)|1),(l + r)/2 + 1,r
// #define int ll
using namespace std;

const int N = 2e5 + 3;

int color[N],vis[N],centroid,f[N];
int tot[N],cnt[N],ans,parent[N];
vector <int> v[N],vec[N];

int dfs(int x,int par,int n,int dl){

	// cout<<x<<" & "<<color[x]<<endl;
	cnt[color[x]] += dl;
	int s = 1;
	for (int to: v[x]){
		if (to == par || vis[to]) continue;
		s += dfs(to,x,n,dl);
	}

	if (centroid == -1 && (par == -1 || s*2 >= n)) 
		centroid = x;
	
	return s;
}
void DFS(int x,int par){
	parent[x] = par;
	for (int to: v[x]){
		if(to==par || vis[to]) continue;
		DFS(to,x);
	}
}
void go(int x,int n,int dep){
	centroid = -1;
	dfs(x,-1,n,1);
	DFS(centroid,-1);

	queue <int> q;
	vector <int> rl;
	q.push(color[centroid]);
	rl.pb(color[centroid]);
	f[color[centroid]] = 1;

	int check = 1,add = 0;
	while (q.size()){
		int c = q.front(); q.pop();

		add++;
		if (tot[c] != cnt[c]) {check = 0; break;}
		
		for (int x: vec[c]){
			int p = parent[x];
			if (p != -1 && color[p] != color[x] && !f[color[p]]){
				f[color[p]] = 1; 
				rl.pb(color[p]);
				q.push(color[p]);
			} 
		}
	}

	for (int x: rl) 
		f[x] = 0;

	if (check) ans=min(ans,add - 1);
	dfs(x,-1,n,-1);
	vis[centroid] = 1;

	for (int to: v[centroid]){
		if(!vis[to]) go(to,n/2,dep + 1);
	}

}
signed main(){
	ios_base::sync_with_stdio(0),cin.tie(NULL),cout.tie(NULL);

    int n,k;
	cin>>n>>k;

	for (int i = 1; i < n; i++){
		int a,b;
		cin>>a>>b;
		v[a].pb(b);
		v[b].pb(a);
	}

	for (int i = 1; i <= n; i++){
		cin >> color[i];
		tot[color[i]]++;
		vec[color[i]].pb(i);
	}

	ans = n;
	go(1,n,1);
	cout<<ans<<endl;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...