Submission #262615

#	Time	Username	Problem	Language	Result	Execution time	Memory
262615		dennisstar	Capital City (JOI20_capital_city)	C++17	100 / 100	1020 ms	40340 KiB

capital_city

#include <bits/stdc++.h>
#define eb emplace_back
using namespace std;
using pii=pair<int, int>;

const int MX = 2e5 + 5;

int N, K, C[MX];
int A, vis[MX], sz[MX];
vector<int> adj[MX], ele[MX];

int P[MX], cnt[MX], use[MX], clh[MX];
void d1(int n, int p) {
	sz[n]=1; cnt[C[n]]++;
	for (auto &i:adj[n])
		if (i!=p&&!vis[i]) d1(i, n), sz[n]+=sz[i];
}

int d2(int n, int p, int c) {
	for (auto &i:adj[n])
		if (i!=p&&!vis[i]&&sz[i]>=c) return d2(i, n, c);
	return n;
}

void d3(int n) {
	vector<pii> V;
	V.eb(n, 0);
	for (int i=0; i<V.size(); i++) {
		int x, y;
		x=V[i].first, y=V[i].second;
		P[x]=y;
		for (auto &j:adj[x])
			if (j!=y&&!vis[j]) V.eb(j, x);
	}
}

void d4(int n) {
	vector<pii> V;
	V.eb(n, 0);
	for (int i=0; i<V.size(); i++) {
		int x, y;
		x=V[i].first, y=V[i].second;
		cnt[C[x]]=0;
		for (auto &j:adj[x])
			if (j!=y&&!vis[j]) V.eb(j, x);
	}
}

void sol(int st) {
	d1(st, 0);
	int c=d2(st, 0, sz[st]/2+1);
	d3(c);

	int fl=0, d=0;
	vector<int> s; s.eb(c);
	for (int i=0; i<s.size(); i++) {
		int x=s[i];
		if (use[x]) continue;
		use[x]=1;
		if (ele[C[x]].size()!=cnt[C[x]]) { fl=1; break; }
		if (!clh[C[x]]) {
			clh[C[x]]=1;
			d++;
			for (auto &j:ele[C[x]]) s.eb(j);
		}
		if (P[x]) s.eb(P[x]);
	}
	if (!fl) A=min(A, d);
	for (auto &i:s) use[i]=0, clh[C[i]]=0;
	d4(st);
	s.clear();

	vis[c]=1;
	for (auto &i:adj[c]) if (!vis[i]) sol(i);
}

int main() {
	cin.tie(0)->sync_with_stdio(0);

	cin>>N>>K;
	for (int i=1, u, v; i<N; i++) cin>>u>>v, adj[u].eb(v), adj[v].eb(u);
	for (int i=1; i<=N; i++) cin>>C[i], ele[C[i]].eb(i);

	A=MX;
	sol(1);

	cout<<A-1<<'\n';
	return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'void d3(int)':
capital_city.cpp:28:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   28 |  for (int i=0; i<V.size(); i++) {
      |                ~^~~~~~~~~
capital_city.cpp: In function 'void d4(int)':
capital_city.cpp:40:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   40 |  for (int i=0; i<V.size(); i++) {
      |                ~^~~~~~~~~
capital_city.cpp: In function 'void sol(int)':
capital_city.cpp:56:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   56 |  for (int i=0; i<s.size(); i++) {
      |                ~^~~~~~~~~
capital_city.cpp:60:23: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   60 |   if (ele[C[x]].size()!=cnt[C[x]]) { fl=1; break; }
      |       ~~~~~~~~~~~~~~~~^~~~~~~~~~~

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