Submission #229769

#	Time	Username	Problem	Language	Result	Execution time	Memory
229769		someone_aa	Capital City (JOI20_capital_city)	C++17	100 / 100	1000 ms	40892 KiB

capital_city

#include <bits/stdc++.h>
#define ll long long
#define pb push_back
#define mp make_pair
#define sz(x) int(x.size())
using namespace std;
const int maxn = 200100;
int n, k;

vector<int>g[maxn], towns[maxn];
int city[maxn], cnt[maxn];

int usize[maxn], uparent[maxn];

int ufind(int node) {
	while(node != uparent[node]) {
		node = uparent[node];
	}
	return node;
}

bool unite(int a, int b) {
	a = ufind(a);
	b = ufind(b);

	if(a == b) return false;
	
	if(usize[a] > usize[b]) {
		uparent[b] = a;
		usize[a] += usize[b];
	}
	else {
		uparent[a] = b;
		usize[b] += usize[a];
	}
	return true;
}

bool deleted[maxn];
int sbt_sz[maxn];
int par[maxn];
int tpar[maxn];

int result;

vector<int>curr_nodes;

void get_sizes(int node, int p) {
	sbt_sz[node] = 1;
	tpar[node] = p;

	for(int i:g[node]) {
		if(deleted[i] || i == p) continue;

		get_sizes(i, node);
		sbt_sz[node] += sbt_sz[i];
	}
}

void get_parents(int node, int p) {
	par[node] = p;
	curr_nodes.pb(node);
	for(int i:g[node]) {
		if(deleted[i] || i == p) continue;
		get_parents(i, node);
	}
}

int get_centroid(int node, int p, int r) {
	for(int i:g[node]) {
		if(i != p && !deleted[i] && 2*sbt_sz[i] >= sbt_sz[r]) {
			return get_centroid(i, node, r);
		}
	}
	return node;
}

void solve(int root) {
	get_sizes(root, -1); // Preprocess and find the subtree sizes of each node
	root = get_centroid(root, root, root); // Find the centroid in this tree
	get_parents(root, root);

	for(int i:curr_nodes) {
		cnt[city[i]]++;
	}

	//cout<<"Centroid: "<<root<<"\n";

	bool check = (cnt[city[root]] == sz(towns[city[root]]));
	int moves = 0;
	if(check) {
		queue<int>q;
		for(int i:towns[city[root]]) {
			q.push(i);
		}

		while(!q.empty()) {
			int curr = q.front();
			q.pop();

			//cout<<"Curr node: "<<curr<<" - "<<par[curr]<<"\n";

			if(unite(city[curr], city[par[curr]])) {
				//cout<<"merge: "<<city[curr]<<" with "<<city[par[curr]]<<"\n";
				moves++;
				check &= (cnt[city[par[curr]]] == sz(towns[city[par[curr]]]));

				if(!check) break;

				for(int i:towns[city[par[curr]]]) {
					q.push(i);
				}
			}
		}
	}

	if(check) result = min(result, moves);

	for(int i:curr_nodes) {
		usize[city[i]] = 1;
		uparent[city[i]] = city[i];
		cnt[city[i]] = 0;
	}

	curr_nodes.clear();
	deleted[root] = true;
	for(int i:g[root]) {
		if(!deleted[i]) {
			solve(i);
		}
	}
}

int main() {
	cin>>n>>k;

	int a, b;
	for(int i=0;i<n-1;i++) {
		cin>>a>>b;
		g[a].pb(b);
		g[b].pb(a);
	}

	for(int i=1;i<=n;i++) {
		cin>>city[i];
		towns[city[i]].pb(i);
	}

	result = INT_MAX;
	for(int i=1;i<=k;i++) {
		usize[i] = 1;
		uparent[i] = i;
	}

	solve(1);
	cout<<result;
}

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