Submission #359853

#	Time	Username	Problem	Language	Result	Execution time	Memory
359853		soroush	Capital City (JOI20_capital_city)	C++14	100 / 100	866 ms	43628 KiB

capital_city

/*
#pragma GCC optimize("O2")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("avx,avx2,sse,sse2,fma")
*/
#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int , int> pii;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int maxn = 2e5 + 100;
const ll mod = 1e9+7;
const ld PI = acos((ld)-1);

#define pb push_back
#define endl '\n'
#define dokme(x) cout << x , exit(0)
#define migmig ios::sync_with_stdio(false),cin.tie(0),cout.tie(0)
#define ms(x , y) memset(x , y , sizeof x)
ll pw(ll a, ll b, ll md = mod){ll res = 1;while(b){if(b&1){res=(a*res)%md;}a=(a*a)%md;b>>=1;}return(res);}

vector < int > adj[maxn], col[maxn];
int n , k , c[maxn];
int sz[maxn] , hide[maxn] , ans = 1e9;
int par[maxn];

void plant(int v , int p = 0){
	sz[v] = 1;
	for(int u : adj[v])if(u ^ p and !hide[u])
		plant(u , v) , sz[v] += sz[u];
}

int cen(int v , int n , int par = 0 , bool found = 0){
	while(!found){
		found = 1;
		for(int u : adj[v]) if(!hide[u] and u != par and sz[u] * 2 > n){
			par = v , v = u , found = 0;
			break;
		}
	}
	return(v);
}

struct sett{
	vector < int > und;
	int cnt[maxn];
	void insert(int x){
		cnt[x] ++;
		if(cnt[x] == 1)und.pb(x);
	}
	void clear(){
		for(int x : und)cnt[x] = 0;
		und.clear();
	}
	int count(int x){
		return(cnt[x]);
	}
}st;

queue < int > colors;

int mark[maxn];
int comp[maxn] , cmp = 0;

void dfs(int v , int p = 0){
	comp[v] = cmp , mark[v] = 0 , par[v] = p;
	for(auto u : adj[v])if(!hide[u]  and u != p)
		par[u] = v , dfs(u , v);
}

int calc(int v){
	cmp = v;
	dfs(v);
	while(colors.size())
		colors.pop();
	colors.push(c[v]);
	st.clear();
	st.insert(c[v]);
	int ans = 0;
	while(!colors.empty()){
		ans ++;
		int color = colors.front();
		colors.pop();
		for(auto u : col[color]){
			if(comp[u] != v)return(1e9);
			if(u != v and !st.count(c[par[u]]))
				st.insert(c[par[u]]) , colors.push(c[par[u]]);
		}
	}
	return(ans - 1);
}

void solve(int v){
	plant(v);
	v = cen(v , sz[v]);
	ans = min(ans , calc(v));
	hide[v] = 1;
	for(int u : adj[v])
		if(!hide[u])
			solve(u);
}

int32_t main(){
	migmig;
	cin >> n >> k;
	if(k > n/2)dokme(0);
	for(int i = 1 ; i < n; i ++){
		int u , v;
		cin >> u >> v;
		adj[u].pb(v);
		adj[v].pb(u);
	}
	for(int i = 1 ; i <= n ; i ++) cin >> c[i] , col[c[i]].pb(i);
	solve(1);
	cout << ans;
	return(0);
}

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