Submission #562805

#	Time	Username	Problem	Language	Result	Execution time	Memory
562805		Mazaalai	수도 (JOI20_capital_city)	C++17	0 / 100	122 ms	25096 KiB

capital_city

#include <bits/stdc++.h>
#define LINE "--------------\n"
#define pb push_back
using namespace std;
const int N = 2e5 + 5;
const int K = 2e5 + 5;
const int INF = 1.07e9;
int n, k;
vector <int> paths[N];
int city[N], pos[N];
vector <int> nums;

void dfs(int pos, int par = 0) {
	nums.pb(city[pos]);
	for (auto el : paths[pos]) {
		if (el != par) dfs(el, pos);
	}
}
int ans[N], ansVal[N];
int cnt[K], l[K], r[K];

int find(int a) {
	if (ans[a] == a) return a;
	return ans[a] = find(ans[a]);
}
void merge(int a, int b) {
	a = find(a);
	b = find(b);
	if (a == b) return;
	ans[b] = a;
	ansVal[a] += ansVal[b]+1;
}
signed main() {
	ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
	// freopen("0.in", "r", stdin);
	// freopen("0.out", "w", stdout);
	cin >> n >> k;
	for (int i = 1; i < n; i++) {
		int a, b; cin >> a >> b;
		paths[a].pb(b);
		paths[b].pb(a);
	}
	for (int i = 1; i <= n; i++) cin >> city[i];
	int leaf = 1;
	for (int i = 1; i <= n; i++) leaf = (paths[i].size() == 1 ? i : leaf);
	nums.pb(-1);
	dfs(leaf);
	stack <int> stk;
	iota(ans+1, ans+1+n, 1);
	// check 0 
	for (int i = 1; i <= k; i++) {
		cnt[i] = 0; l[i] = n+1, r[i] = 0;
	}
	for (int i = 1; i <= n; i++) {
		cnt[city[nums[i]]]++;
		l[city[nums[i]]] = min(l[city[nums[i]]], i);
		r[city[nums[i]]] = i;
	}
	int ans = INF;
	for (int i = 1; i <= k; i++) {
		if (cnt[i] == 0) continue;
		if (cnt[i] == r[i] - l[i] + 1) ans = 0;
	}
	for (int i = 1; i <= n; i++) {
		int x = nums[i];
		while (!stk.empty() && pos[city[x]] != 0 && stk.top() >= pos[city[x]])  {
			merge(x, nums[stk.top()]);
			stk.pop();
		}
		stk.push(i);
		pos[city[x]] = i;
		if (r[city[x]] == i) ans = min(ans, ansVal[find(city[x])]);
		// cout << LINE;
		// cout << x << '\n';
		// for (int j = 1; j <= k; j++) cout << find(j) << " \n"[j==k];
		// for (int j = 1; j <= k; j++) cout << ansVal[find(j)] << " \n"[j==k];
	}
	cout << ans << '\n';
}

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