Submission #609708

#	Time	Username	Problem	Language	Result	Execution time	Memory
609708		penguinhacker	Capital City (JOI20_capital_city)	C++17	100 / 100	1032 ms	163708 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;
 
#define ll long long
#define ar array
 
const int mxN=2e5, MX=1e6;
int n, k, c[mxN], n2;
int d[mxN], s[mxN], p[mxN], hd[mxN], tin[mxN], tin2[mxN];
vector<int> oc[mxN], adj1[mxN], nodes;
vector<int> adj2[MX], adj3[MX], st;
int who[MX], cnt[MX];
bool vis[MX], out[MX];
 
void dfs1(int u=0) {
	s[u]=1;
	for (int& v : adj1[u]) {
		adj1[v].erase(find(adj1[v].begin(), adj1[v].end(), u));
		d[v]=d[u]+1, p[v]=u;
		dfs1(v);
		s[u]+=s[v];
		if (s[v]>s[adj1[u][0]])
			swap(v, adj1[u][0]);
	}
}

void dfs2(int u=0, int h=0) {
	static int timer=0;
	hd[u]=h, tin[u]=timer++, tin2[tin[u]]=u;
	if (adj1[u].empty())
		return;
	dfs2(adj1[u][0], h);
	for (int i=1; i<adj1[u].size(); ++i)
		dfs2(adj1[u][i], adj1[u][i]);
}

int lca(int u, int v) {
	for (; hd[u]!=hd[v]; v=p[hd[v]])
		if (d[hd[u]]>d[hd[v]])
			swap(u, v);
	return d[u]<d[v]?u:v;
}

void ae(int u, int v) {
	adj2[u].push_back(v);
	adj3[v].push_back(u);
}

void bld(int u=1, int l=0, int r=n-1) {
	if (l==r) {
		ae(u+k, c[tin2[l]]);
		n2=max(n2, u+k+1);
		return;
	}
	int mid=(l+r)/2;
	ae(u+k, 2*u+k);
	ae(u+k, 2*u+1+k);
	bld(2*u, l, mid);
	bld(2*u+1, mid+1, r);
}

void qry(int ql, int qr, int u=1, int l=0, int r=n-1) {
	if (l>qr||r<ql)
		return;
	if (ql<=l&&r<=qr) {
		nodes.push_back(u+k);
		return;
	}
	int mid=(l+r)/2;
	qry(ql, qr, 2*u, l, mid);
	qry(ql, qr, 2*u+1, mid+1, r);
}

vector<int> gt(int u, int v) {
	assert(hd[u]==hd[v]&&tin[u]<=tin[v]);
	nodes.clear();
	qry(tin[u], tin[v]);
	return nodes;
}

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

void dfs4(int u, int rt) {
	vis[u]=0, who[u]=rt;
	cnt[rt]+=u<k;
	for (int v : adj3[u])
		if (vis[v])
			dfs4(v, rt);
}
 
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();
	dfs2();
	bld();
	for (int i=0; i<n; ++i) {
		ae(i+n2, c[i]);
		if (i!=hd[i])
			ae(i+n2, p[i]+n2);
	}
	for (int i=0; i<k; ++i) {
		int l=oc[i][0];
		for (int j=1; j<oc[i].size(); ++j)
			l=lca(l, oc[i][j]);
		for (int u : oc[i]) {
			for (; hd[u]!=hd[l]; u=p[hd[u]])
				ae(i, n2+u);
			for (int j : gt(l, u))
				ae(i, j);
		}
	}
	n2+=n;
	for (int i=0; i<n2; ++i)
		if (!vis[i])
			dfs3(i);
	reverse(st.begin(), st.end());
	for (int i : st)
		if (vis[i])
			dfs4(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 'void dfs2(int, int)':
capital_city.cpp:33:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   33 |  for (int i=1; i<adj1[u].size(); ++i)
      |                ~^~~~~~~~~~~~~~~
capital_city.cpp: In function 'int main()':
capital_city.cpp:121:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  121 |   for (int j=1; j<oc[i].size(); ++j)
      |                 ~^~~~~~~~~~~~~

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