Submission #366042

#	Time	Username	Problem	Language	Result	Execution time	Memory
366042		negar_a	Capital City (JOI20_capital_city)	C++14	100 / 100	879 ms	42340 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef long double ld;
#define pb push_back
#define mp make_pair
#define pii pair <int, int>
#define fast_io ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define F first
#define S second

const int maxn = 2e5 + 5;
vector <int> adj[maxn], vec[maxn];
int sz[maxn], par[maxn], c[maxn];
int ans;
bool cent[maxn], mark[maxn], com[maxn], col[maxn];
vector <int> ver;

void dfs(int v, int p = -1){
	sz[v] = 1;
	com[v] = true;
	ver.pb(v);
	for(int u : adj[v]){
		if(!cent[u] && u != p){
			dfs(u, v);
			sz[v] += sz[u];
		}
	}
}
int find_centroid(int v, int s, int p = -1){
	for(int u : adj[v]){
		if(!cent[u] && u != p && sz[u] > s / 2){
			return find_centroid(u, s, v);
		}
	}
	return v;
}
void dfs2(int v, int p = -1){
	for(int u : adj[v]){
		if(!cent[u] && u != p){
			par[u] = v;
			dfs2(u, v);
		}
	}
}

void solve(int v){
	ver.clear();
	dfs(v);
	int cn = find_centroid(v, sz[v]);
//	cout << "cen : " << cn << endl; 
	
	par[cn] = -1;
	dfs2(cn);
	queue <int> q;
	vector <int> vc;
	q.push(cn);
	
	bool flag = true; 
	while(q.size() && flag){
		int x = q.front();
		q.pop();
		if(!col[c[x]]){
			col[c[x]] = true;
			vc.pb(c[x]);
			for(int u : vec[c[x]]){
				if(!com[u]){
					flag = false;
					break;
				}
				if(!mark[u]){
					q.push(u);
					mark[u] = true;
				}
			}	
		}
		if(par[x] != -1){
			if(!com[par[x]]){
				flag = false;
				break;
			}
			if(!mark[par[x]]){
				q.push(par[x]);
				mark[par[x]] = true;
			}
		}
	}
	if(flag)
		ans = min(ans, (int)vc.size() - 1);
		
	//clear
	for(int u : ver){
		col[c[u]] = false;
		mark[u] = sz[u] = par[u] = com[u] = 0;
	}
		
	cent[cn] = true;
	for(int u : adj[cn]){
		if(!cent[u]){
			solve(u);
		}
	}
}

int main(){
	fast_io;
	int n, k;
	cin >> n >> k;
	for(int i = 0; i < n - 1; i ++){
		int x, y;
		cin >> x >> y;
		x --; y --;
		adj[x].pb(y);
		adj[y].pb(x);
	}
	for(int i = 0; i < n; i ++){
		cin >> c[i];
		vec[c[i]].pb(i);
	}
	ans = k - 1;
	solve(0);
	cout << ans;
	
	return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'void solve(int)':
capital_city.cpp:95:19: warning: suggest parentheses around assignment used as truth value [-Wparentheses]
   95 |   mark[u] = sz[u] = par[u] = com[u] = 0;
      |             ~~~~~~^~~~~~~~~~~~~~~~~~~~~

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