Submission #220221

#	Time	Username	Problem	Language	Result	Execution time	Memory
220221		atoiz	Capital City (JOI20_capital_city)	C++14	100 / 100	2491 ms	427176 KiB

capital_city

/*input
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
*/

#include <iostream>
#include <vector>
#include <cstdio>
#include <algorithm>
#include <cmath>
#include <stack>
#include <cassert>

using namespace std;

const int MAXN = 200007, LOG = __lg(MAXN) + 1, MAXV = MAXN * LOG;
int N, lg, anc[LOG][MAXN], dep[MAXN], city[MAXN], K;
vector<int> towns[MAXN];
vector<int> tree[MAXN], G[MAXV];

int low[MAXV], vis[MAXV], col[MAXV], T = 0, ans = MAXN;
stack<int> st;
bool in_comp[MAXV];

void dfs(int u, int p)
{
	for (int v : tree[u]) if (v != p) {
		anc[0][v] = u;
		dep[v] = dep[u] + 1;
		dfs(v, u);
	}
}

int get_lca(int u, int v)
{
	if (dep[u] < dep[v]) swap(u, v);
	for (int j = lg - 1; j >= 0; --j) if (dep[u] - (1 << j) >= dep[v]) u = anc[j][u];
	if (u == v) return u;
	for (int j = lg - 1; j >= 0; --j) if (anc[j][u] != anc[j][v]) u = anc[j][u], v = anc[j][v];
	return anc[0][u];
}

void dfs2(int u)
{
	st.push(u);
	low[u] = vis[u] = ++T; col[u] = 1;

	for (int v : G[u]) {
		if (col[v] == 2);
		else if (col[v] == 1) {
			low[u] = min(low[u], vis[v]);
		} else {
			// cout << (u - 1) / N << ' ' << (u - (u - 1) / N * N) << ": " << (v - 1) / N << ' ' << (v - (v - 1) / N * N) << endl;
			dfs2(v);
			low[u] = min(low[u], low[v]);
		}
	}

	if (low[u] == vis[u]) {
		int cost = 0;
		vector<int> vec;
		// cout << "comp: ";
		for (int w = st.top();; w = st.top()) {
			st.pop();
			vec.push_back(w);
			col[w] = 2;
			// cout << '(' << (w - 1) / N << ' ' << (w - (w - 1) / N * N) << ") ";
			if (w == u) break;
		}

		bool valid = true;
		for (int w : vec) in_comp[w] = 1;
		for (int w : vec) for (int v : G[w]) if (!in_comp[v]) { valid = false; break; }
		for (int w : vec) in_comp[w] = 0;

		if (valid) {
			int cost = 0;
			for (int w : vec) if (w <= N && !in_comp[city[w]]) {
				cost += (in_comp[city[w]] = 1);
			}
			ans = min(ans, cost);
			for (int w : vec) if (w <= N && in_comp[city[w]]) {
				in_comp[city[w]] = 0;
			}
		}
		// cout << valid << endl;
	}
}

int main()
{
	ios_base::sync_with_stdio(0); cin.tie(0);
	cin >> N >> K;
	lg = __lg(N) + 1;
	int V = lg * N;
	for (int i = 0; i < N - 1; ++i) {
		int u, v;
		cin >> u >> v;
		tree[u].push_back(v);
		tree[v].push_back(u);
	}
	dep[1] = 1, dfs(1, 0);
	for (int j = 1; j < lg; ++j) for (int u = 1; u <= N; ++u) {
		anc[j][u] = anc[j - 1][anc[j - 1][u]];
		if (dep[u] >= (1 << j)) {
			G[j * N + u].push_back((j - 1) * N + u);
			G[j * N + u].push_back((j - 1) * N + anc[j - 1][u]);
		}
	}

	for (int u = 1; u <= N; ++u) {
		int c;
		cin >> c;
		towns[c].push_back(u);
		city[u] = c;
	}
	for (int c = 1; c <= K; ++c) {
		int lca = towns[c][0];
		for (int v : towns[c]) lca = get_lca(lca, v);
		lca = anc[0][lca];

		for (int v : towns[c]) {
			int w = v;
			for (int j = lg - 1; j >= 0; --j) if (dep[w] - (1 << j) >= dep[lca]) {
				G[v].push_back(j * N + w);
				w = anc[j][w];
			}

			if (v != towns[c][0]) {
				G[v].push_back(towns[c][0]);
				G[towns[c][0]].push_back(v);
			}
		}
	}

	// cout << V << endl;
	for (int j = 0; j < lg; ++j) for (int u = 1; u <= N; ++u) {
		if (dep[u] >= (1 << j) && !col[j * N + u]) {
			dfs2(j * N + u);
		}
	}
	cout << ans - 1 << endl;
}

Compilation message (stderr)

capital_city.cpp: In function 'void dfs2(int)':
capital_city.cpp:71:7: warning: unused variable 'cost' [-Wunused-variable]
   int cost = 0;
       ^~~~
capital_city.cpp: In function 'int main()':
capital_city.cpp:106:6: warning: unused variable 'V' [-Wunused-variable]
  int V = lg * N;
      ^

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