Submission #911473

#	Time	Username	Problem	Language	Result	Execution time	Memory
911473		OAleksa	Capital City (JOI20_capital_city)	C++14	100 / 100	315 ms	87360 KiB

capital_city

#include <bits/stdc++.h>
//ako ovaj vaso daso misli da me pobedjuje.....
using namespace std;
#define int long long
#define f first
#define s second
const int N = 2e5 + 69;
const int K = 18;
int n, a[N], k, lca[N], vis[N];
int up[N][K], tin[N], tout[N], timer;
vector<int> g[N], t[N], rev[N], c[N];
vector<int> ord;
bool anc(int a, int b) {
	return tin[a] <= tin[b] && tout[a] >= tout[b];
}
int LCA(int a, int b) {
	if (anc(a, b))
		return a;
	else if (anc(b, a))
		return b;
	for (int i = K - 1;i >= 0;i--) {
		if (!anc(up[a][i], b) && up[a][i] > 0)
			a = up[a][i];
	}
	return up[a][0];
}
void dfs(int v, int p) {
	up[v][0] = p;
	tin[v] = ++timer;
	for (int i = 1;i < K;i++)
		up[v][i] = up[up[v][i - 1]][i - 1];
	for (auto u : t[v]) {
		if (u == p)
			continue;
		dfs(u, v);
	}
	tout[v] = timer;
}
void dfs1(int v) {
	if (vis[v])
		return;
	vis[v] = 1;
	for (auto u : g[v]) {
		dfs1(u);
	}
	ord.push_back(v);
}
int s[N], comp, sz[N];
vector<int> tmp;
void dfs2(int v) {
	if (vis[v])
		return;
	vis[v] = 1;
	tmp.push_back(v);
	for (auto u : rev[v]) 
		dfs2(u);
}
signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  int tt = 1;
  //cin >> tt;
  while (tt--) {
  	cin >> n >> k;
  	for (int i = 1;i <= n - 1;i++) {
  		int a, b;
  		cin >> a >> b;
  		t[a].push_back(b);
  		t[b].push_back(a);
  	}
  	for (int i = 1;i <= n;i++) {
  		cin >> a[i];
  		c[a[i]].push_back(i);
  	}
  	dfs(1, 0);
  	for (int i = 1;i <= k;i++) {
  		assert(c[i].size() > 0);
  		lca[i] = c[i][0];
  		for (int j = 1;j < (int)c[i].size();j++)
  			lca[i] = LCA(lca[i], c[i][j]);
  	}
  	for (int i = 1;i <= n;i++) {
  		if (i == lca[a[i]])
  			continue;
  		if (a[i] != a[up[i][0]]) {
  			g[a[i]].push_back(a[up[i][0]]);
  			rev[a[up[i][0]]].push_back(a[i]);
  		}
  	}
  	for (int i = 1;i <= n;i++) {
  		sort(g[i].begin(), g[i].end());
  		g[i].erase(unique(g[i].begin(), g[i].end()), g[i].end());
  		sort(rev[i].begin(), rev[i].end());
  		rev[i].erase(unique(rev[i].begin(), rev[i].end()), rev[i].end());
  	}
  	int ans = 1e9;
  	for (int i = 1;i <= k;i++) {
  		if (!vis[i])
  			dfs1(i);
  	}
  	for (int i = 1;i <= n;i++)
  		vis[i] = 0;
  	reverse(ord.begin(), ord.end());
  	for (auto u : ord) {
  		if (!vis[u]) {
  			++comp;
  			tmp.clear();
  			dfs2(u);
  			sz[comp] = (int)tmp.size();
  			for (auto x : tmp)
  				s[x] = comp;
  		}
  	}
  	for (int i = 1;i <= k;i++) {
  		int ok = 1;
  		for (auto u : g[i]) {
  			if (s[u] != s[i])
  				ok = 0;
  		}
  		if (!ok)
  			sz[s[i]] = 1e9;
  	}
  	for (int i = 1;i <= comp;i++)
  		ans = min(ans, sz[i] - 1);
  	cout << ans;
	}
  return 0;
}

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