Submission #427966

#	Time	Username	Problem	Language	Result	Execution time	Memory
427966		milleniumEeee	Capital City (JOI20_capital_city)	C++17	11 / 100	3049 ms	60876 KiB

capital_city

#include <bits/stdc++.h>
#define fr first
#define sc second
#define pii pair<int, int>
#define pb push_back
#define szof(s) (int)s.size()
#define all(s) s.begin(), s.end()
#define fastInp ios_base::sync_with_stdio(0); cin.tie(0);
#define ll long long
template<class T>void chmax(T &a, T b){if (a < b)a = b;}
template<class T>void chmin(T &a, T b){if (b < a)a = b;}

using namespace std;

const int MAXN  = (int)2e5 + 5;
const int INF = 1e9;
const int L = 20;

vector <int> g[MAXN];
int c[MAXN];
int pr[MAXN][L + 1];
int tin[MAXN], tout[MAXN], tiktak = 0;

vector <int> color_comp[MAXN];

void precalc(int v, int par) {
	tin[v] = ++tiktak;
	pr[v][0] = par;
	for (int lv = 1; lv <= L; lv++) {
		pr[v][lv] = pr[pr[v][lv - 1]][lv - 1];
	}
	for (int to : g[v]) {
		if (to != par) {
			precalc(to, v);
		}
	}
	tout[v] = tiktak;
}

bool father(int a, int b) {
	return tin[a] <= tin[b] && tout[b] <= tout[a];
}

int get_lca(int a, int b) {
	if (father(a, b)) {
		return a;
	}
	if (father(b, a)) {
		return b;
	}
	for (int lv = L; lv >= 0; lv--) {
		if (!father(pr[a][lv], b)) {
			a = pr[a][lv];
		}
	}
	return pr[a][0];
}

int get_lca(vector <int> &vec) {
	int lca = vec[0];
	for (int i = 1; i < szof(vec); i++) {
		lca = get_lca(lca, vec[i]);
	}
	return lca;
}

int cur[MAXN];

int used[MAXN], used_id = 1;

void up_dfs(vector <int> &vec, int v, int root) {
	if (used[v] == used_id) {
		return;
	}
	used[v] = used_id;
	vec.pb(v);
	if (v != root) {
		up_dfs(vec, pr[v][0], root);		
	}	
}

void relax(vector <int> &comp) {
	used_id++;
	vector <int> res;
	int lca = get_lca(comp);
	for (int v : comp) {
		up_dfs(res, v, lca);
	}
	comp = res;
}

signed main() {
	fastInp;
	int n, k;
	cin >> n >> k;
	for (int i = 1; i < n; i++) {
		int u, v;
		cin >> u >> v;
		g[u].pb(v);
		g[v].pb(u);
	}
	set <int> colors;
	for (int i = 1; i <= n; i++) {
		cin >> c[i];
		colors.insert(c[i]);
		color_comp[c[i]].pb(i);
	}
	precalc(1, 1);
	int ans = INF;
	for (int cur_color : colors) {
		for (int i = 1; i <= n; i++) {
			cur[i] = c[i];
		}
		vector <int> comp;
		for (int i = 1; i <= n; i++) {
			if (c[i] == cur_color) {
				comp.pb(i);
			}
		}
		int cur_ans = 0;
		while (1) {
			relax(comp);
			bool ok = 1;
			set <int> new_colors;
			for (int v : comp) {
				ok &= (cur[v] == cur_color);
				if (cur[v] != cur_color) {
					new_colors.insert(cur[v]);
				}
			}
			if (ok) {
				break;
			}
			else {
				for (int el : new_colors) {
					cur_ans++;
					for (int v : color_comp[el]) {
						comp.pb(v);
						cur[v] = cur_color;
					}
				}
				continue;
			}
		}
		chmin(ans, cur_ans);
	}
	cout << ans << endl;
}

/*
12 4
7 9
1 3
4 6
2 4
10 12
1 2
2 10
11 1
2 8
5 3
6 7
3
1
1
2
4
3
3
2
2
3
4
4

*/

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