Submission #414341

#TimeUsernameProblemLanguageResultExecution timeMemory
414341ritul_kr_singhCapital City (JOI20_capital_city)C++17
100 / 100
2257 ms36972 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;
 
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){
	s[p] -= s[u], s[u] = sz;
	for(int v : g[u]) if(!r[v] && s[v] > sz/2) return find(v, u, sz);
	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;
	queue<int> q;
	for(int i : h[color[c]]){
		if(res == MAXN) break;
		if(par[i] < 0) res = MAXN;
		else q.push(i);
	}
 
	while(!q.empty() && res < MAXN){
		int u = q.front(); q.pop();
		if(!vis[color[par[u]]]){
			++res;
			for(int v : h[color[par[u]]]){
				if(par[v] >= 0 && res < MAXN) q.push(v);
				else res = MAXN;
			}
			if(res < MAXN) vis[color[par[u]]] = 1;
		}
	}
	ans = min(ans, res);
 
	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...