Submission #414312

#TimeUsernameProblemLanguageResultExecution timeMemory
414312ritul_kr_singhCapital City (JOI20_capital_city)C++17
11 / 100
3069 ms37316 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;
	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;
	vector<int> rollback = {color[c]};
	vis[color[c]] = 1;
	queue<int> q;
	for(int i : h[color[c]]){
		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;
			rollback.push_back(color[par[u]]);
			vis[color[par[u]]] = 1;
			for(int v : h[color[par[u]]]){
				if(par[v] >= 0 && res < MAXN) q.push(v);
				else res = MAXN;
			}
		}
	}
	for(int i : rollback) vis[i] = 0;
	vector<int> ().swap(rollback);
	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...