Submission #537790

#	Time	Username	Problem	Language	Result	Execution time	Memory
537790		8e7	Capital City (JOI20_capital_city)	C++17	100 / 100	972 ms	45704 KiB

capital_city

//Challenge: Accepted
#include <bits/stdc++.h>
using namespace std;
#ifdef zisk
void debug(){cout << endl;}
template<class T, class ... U> void debug(T a, U ... b){cout << a << " ", debug(b...);}
template<class T> void pary(T l, T r) {
	while (l != r) cout << *l << " ", l++;
	cout << endl;
}
#else
#define debug(...) 0
#define pary(...) 0
#endif
#define ll long long
#define maxn 200005
#define pii pair<int, int>
#define ff first
#define ss second
#define io ios_base::sync_with_stdio(0);cin.tie(0);
vector<int> adj[maxn], pos[maxn];
int c[maxn], siz[maxn];
int pa[maxn];
bool vis[maxn], col[maxn], mark[maxn], done[maxn], ins[maxn];
void dfs(int n, int par, vector<int> &vec) {
	vec.push_back(n);
	pa[n] = par;
	siz[n] = 1;
	for (int v:adj[n]) {
		if (v != par && !mark[v]) {
			dfs(v, n, vec);
			siz[n] += siz[v];
		}
	}
}
int getcent(int n, int par, int tot) {
	for (int v:adj[n]) {
		if (v != par && !mark[v] && siz[v] * 2 >= tot) return getcent(v, n, tot);
	}
	return n;
}
int ans = 1<<30;
void solve(vector<int> a) {
	if (a.size() == 0) return;
	pary(a.begin(), a.end());
	vector<int> tmp;
	dfs(a[0], 0, tmp);	
	int root = getcent(a[0], 0, siz[a[0]]);
	dfs(root, 0, tmp);	
	if (!done[c[root]]) {
		for (int i:a) col[c[i]] = 0, vis[i] = 0, ins[i] = 1;	
		bool check = 1;

		queue<int> que;	
		vis[root] = 1;
		col[c[root]] = 1;
		int cnt = 0;
		for (int v:pos[c[root]]) {
			if (!ins[v]) {
				check = 0;
				break;
			}
			que.push(v);
		}
		while (que.size()) {
			if (!check) break;
			int cur = que.front();que.pop();
			while (!vis[cur]) {
				vis[cur] = 1;
				if (!col[c[cur]]) {
					cnt++;
					col[c[cur]] = 1;	
					for (int v:pos[c[cur]]) {
						if (!ins[v]) {
							check = 0;
							break;
						}
						que.push(v);
					}
				}
				cur = pa[cur];
			}
		}
		if (check) ans = min(ans, cnt);
		done[c[root]] = 1;
		for (int i:a) ins[i] = 0;
	}
	mark[root] = 1;
	for (int v:adj[root]) {
		if (mark[v]) continue;
		vector<int> rec;
		dfs(v, root, rec);
		solve(rec);
	}
}
int main() {
	io
	int n, K;
	cin >> n >> K;
	for (int i = 0;i < n - 1;i++) {
		int u, v;
		cin >> u >> v;
		adj[u].push_back(v);
		adj[v].push_back(u);	
	}
	vector<int> ini;	
	for (int i = 1;i <= n;i++) {
		cin >> c[i];
		pos[c[i]].push_back(i);
		ini.push_back(i);
	}
	solve(ini);	
	cout << ans << endl;
}

Compilation message (stderr)

capital_city.cpp: In function 'void solve(std::vector<int>)':
capital_city.cpp:13:19: warning: statement has no effect [-Wunused-value]
   13 | #define pary(...) 0
      |                   ^
capital_city.cpp:45:2: note: in expansion of macro 'pary'
   45 |  pary(a.begin(), a.end());
      |  ^~~~

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