Submission #1016276

#	Time	Username	Problem	Language	Result	Execution time	Memory
1016276		THXuan	Capital City (JOI20_capital_city)	C++14	100 / 100	270 ms	42064 KiB

capital_city

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>
#include <set>
#include <map>
#define INF 1e9
using namespace std;
typedef long long ll;

int n = 0, k = 0;
int c[200005], dp[200005], root[200005], p[200005];
bool visited[200005], processed[200005];
vector<int>adj[200005], v[200005];
queue<int>q;

int get_size(int s, int e) {
	dp[s] = 1;
	for (auto u : adj[s]) {
		if (u == e || visited[u])continue;
		get_size(u, s);
		dp[s] += dp[u];
	}
	return dp[s];
}

int get_centroid(int s, int e, int n) {
	for (auto u : adj[s]) {
		if (u == e || visited[u])continue;
		if (dp[u] * 2 > n) return get_centroid(u, s, n);
	}
	return s;
}

void dfs(int s, int e = 0) {
	p[s] = e;
	processed[c[s]] = 0;
	root[s] = (e == 0 ? s : root[e]);
	for (int u : adj[s]) {
		if (u != e && !visited[u]) {
			dfs(u, s);
		}
	}
}

int centroid_decomp(int node) {
	int centroid = get_centroid(node, 0, get_size(node, 0));
	visited[centroid] = true;
	dfs(centroid);
	q.push(c[centroid]);
	int ans = 0;
	while (q.size()) {
		int col = q.front();q.pop();
		if (processed[col]) continue;
		processed[col] = true;
		for (int u : v[col]) {
			if (root[u] != centroid) {
				ans = k;
				break;
			}
			if (p[u] != 0) q.push(c[p[u]]);
		}
		ans++;
	}
	//cout << centroid << " " << ans << "\n";
	for (int u : adj[centroid]) {
		if (!visited[u]) ans = min(ans, centroid_decomp(u));
	}
	return ans;
}

void solve()
{
	cin >> n >> k;
	for (int i = 1; i <= n - 1; i++) {
		int u, v; cin >> u >> v;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
	for (int i = 1; i <= n; i++) { cin >> c[i]; v[c[i]].push_back(i); }
	cout << centroid_decomp(1) - 1 << "\n";
}

int main()
{
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	int t = 1;// cin>>t;
	while (t--) solve();
	return 0;
}

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