제출 #414355

#제출 시각아이디문제언어결과실행 시간메모리
414355ritul_kr_singhCapital City (JOI20_capital_city)C++17
100 / 100
708 ms34260 KiB
#include <bits/stdc++.h>
using namespace std;
// #define int long long
#define sp << ' ' <<
#define nl << '\n'
 
const int MAXN = 2e5;
 
int n, k, color[MAXN], par[MAXN], s[MAXN], ans = MAXN;
vector<int> g[MAXN], h[MAXN];
bitset<MAXN> r, vis, bad;
 
void calcSize(int u){
	s[u] = 1;
	for(int v : g[u]) if(!s[v]) calcSize(v), s[u] += s[v];
}
 
int find(int u, int p, int sz){
	bool f = 1;
	while(f){
		f = 0;
		s[p] -= s[u], s[u] = sz;
		for(int v : g[u]){
			if(!r[v] && s[v] > sz/2){
				p = u, u = v;
				f = 1;
				break;
			}
		}
	}
	return u;
}
 
void dfs(int u, int p, bool val){
	par[u] = val ? p : -1;
	if(val) vis[color[u]] = 0;
	for(int v : g[u]) if(v != p && !r[v]) dfs(v, u, val);
}
 
void decompose(int c){
	r[c = find(c, c, s[c])] = 1;
	dfs(c, c, 1);
 
	int res = 0;
	vis[color[c]] = 1;
	vector<int> st;
	for(int i : h[color[c]]){
		if(res == MAXN) break;
		if(par[i] < 0) res = MAXN;
		else st.push_back(i);
	}
 
	while(!st.empty() && res < MAXN){
		int u = st.back(); st.pop_back();
		if(!vis[color[par[u]]]){
			if(bad[color[par[u]]]){
				res = MAXN;
				break;
			}
			++res;
			for(int v : h[color[par[u]]]){
				if(par[v] >= 0 && res < MAXN) st.push_back(v);
				else res = MAXN;
			}
			if(res < MAXN) vis[color[par[u]]] = 1;
			else bad[color[par[u]]] = 1;
		}
	}
	ans = min(ans, res);

	bad[color[c]] = 1; 
	dfs(c, c, 0);
	for(int v : g[c]) if(!r[v]) decompose(v);
}
 
signed main(){
	cin.tie(0)->sync_with_stdio(0);
	int x, y; cin >> n >> k;
	for(int i=1; i<n; ++i){
		cin >> x >> y; --x, --y;
		g[x].push_back(y);
		g[y].push_back(x);
	}
 
	for(int i=0; i<n; ++i){
		cin >> color[i]; --color[i];
		h[color[i]].push_back(i);
	}
 
	calcSize(0);
	decompose(0);
 
	cout << ans;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...