Submission #562811

#	Time	Username	Problem	Language	Result	Execution time	Memory
562811		Mazaalai	Capital City (JOI20_capital_city)	C++17	30 / 100	146 ms	27832 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 rs[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);
	rs[a] = rs[b] = max(rs[a], rs[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[nums[i]]++;
		l[nums[i]] = min(l[nums[i]], i);
		r[nums[i]] = i;
	}
	int res = INF;
	for (int i = 1; i <= k; i++) {
		rs[i] = r[i];
		if (cnt[i] == 0) continue;
		if (cnt[i] > r[i] - l[i]) res = 0;
	}
	for (int i = 1; i <= n; i++) {
		int x = nums[i];
		while (!stk.empty() && pos[x] != 0 && stk.top() >= pos[x])  {
			merge(nums[i], nums[stk.top()]);
			stk.pop();
		}
		stk.push(i);
		pos[x] = i;
		// 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 << rs[find(j)] << " \n"[j==k];
		// for (int j = 1; j <= k; j++) cout << ansVal[find(j)] << " \n"[j==k];

		if (i == rs[find(x)]) {
			// cout << "HERE\n";
			res = min(res, ansVal[find(x)]);
		}
	}
	cout << res << '\n';
}

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