Submission #595745

#	Time	Username	Problem	Language	Result	Execution time	Memory
595745		penguinhacker	Capital City (JOI20_capital_city)	C++17	11 / 100	2556 ms	524288 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;

#define ll long long
#define ar array

const int mxN=2e5;
int n, k, c[mxN], d[mxN], anc[mxN][18], tin[mxN], t, who[20*mxN], cnt[20*mxN];
vector<int> oc[mxN], adj1[mxN], adj2[20*mxN], adj3[20*mxN], st;
bool vis[20*mxN], out[20*mxN];

void dfs1(int u=0) {
	tin[u]=t++;
	for (int i=1; i<18; ++i)
		anc[u][i]=anc[anc[u][i-1]][i-1];
	for (int v : adj1[u])
		if (v!=anc[u][0]) {
			d[v]=d[u]+1, anc[v][0]=u;
			dfs1(v);
		}
}

int lca(int u, int v) {
	if (d[u]>d[v])
		swap(u, v);
	for (int i=17; ~i; --i)
		if (d[v]-(1<<i)>=d[u])
			v=anc[v][i];
	if (u==v)
		return u;
	for (int i=17; ~i; --i)
		if (anc[u][i]!=anc[v][i])
			u=anc[u][i], v=anc[v][i];
	return anc[u][0];
}

void dfs2(int u) {
	vis[u]=1;
	for (int v : adj2[u])
		if (!vis[v])
			dfs2(v);
	st.push_back(u);
}

void dfs3(int u, int rt) {
	vis[u]=0, who[u]=rt;
	cnt[rt]+=u<k;
	for (int v : adj3[u])
		if (vis[v])
			dfs3(v, rt);
}

int cur_ans;

void dfs4(int u) {
	vis[u]=1;
	cur_ans+=u<k;
	for (int v : adj2[u])
		if (!vis[v])
			dfs4(v);
}

void ae(int u, int v) {
	//if (u<20) cout << u << " " << v << endl;
	adj2[u].push_back(v);
	adj3[v].push_back(u);
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cin >> n >> k;
	for (int i=1; i<n; ++i) {
		int u, v;
		cin >> u >> v, --u, --v;
		adj1[u].push_back(v);
		adj1[v].push_back(u);
	}
	for (int i=0; i<n; ++i) {
		cin >> c[i], --c[i];
		oc[c[i]].push_back(i);
	}
	dfs1();
	for (int i=0; i<k; ++i) {
		sort(oc[i].begin(), oc[i].end(), [&](int a, int b) { return tin[a]<tin[b]; } );
		for (int j=1; j<oc[i].size(); ++j) {
			int u=oc[i][j-1], v=oc[i][j];
			int uv=lca(u, v);
			ae(i, c[uv]);
			for (int rep=0; rep<2; ++rep) {
				for (int j=17; ~j; --j)
					if (d[u]-(1<<j)>=d[uv]) {
						ae(i, j*n+u+k);
						u=anc[u][j];
					}
				swap(u, v);
			}
		}
	}
	for (int i=0; i<n; ++i) {
		ae(i+k, c[i]);
		for (int j=1; j<18; ++j) {
			ae(j*n+i+k, (j-1)*n+i+k);
			ae(j*n+i+k, (j-1)*n+anc[i][j-1]+k);
		}
	}
	for (int i=0; i<k; ++i)
		if (!vis[i])
			dfs2(i);
	reverse(st.begin(), st.end());
	for (int i : st)
		if (vis[i])
			dfs3(i, i);
	for (int i : st)
		for (int j : adj2[i])
			if (who[i]!=who[j])
				out[who[i]]=1;
	int ans=69696969;
	for (int i=0; i<k; ++i)
		if (!out[who[i]])
			ans=min(ans, cnt[who[i]]-1);
	cout << ans;
	return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'int main()':
capital_city.cpp:86:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   86 |   for (int j=1; j<oc[i].size(); ++j) {
      |                 ~^~~~~~~~~~~~~

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