Submission #414938

#	Time	Username	Problem	Language	Result	Execution time	Memory
414938		Pro_ktmr	Capital City (JOI20_capital_city)	C++17	100 / 100	2840 ms	97124 KiB

capital_city

#pragma GCC target ("avx")
#pragma GCC optimize ("unroll-loops")
#pragma GCC optimize ("O3")

#include"bits/stdc++.h"
#include<unordered_set>
#include<unordered_map>
#include<random>
using namespace std;
typedef long long ll;
const ll MOD = (ll)(1e9+7);
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
int dx[4]={ 1,0,-1,0 };
int dy[4]={ 0,1,0,-1 };

int N, K;
vector<int> e[200000];
int C[200000];

int G[200000] ={};
bool used[200000] ={};

int sub[200000];
int par[200000];
vector<int> chi[200000];
bool ok[200000];
int dfs(int n, int p, int N){
	sub[n] = 1;
	par[n] = p;
	chi[n].clear();
	ok[n] = true;
	rep(i, e[n].size()){
		if(e[n][i] == p || used[e[n][i]]) continue;
		chi[n].pb(e[n][i]);
		int ret = dfs(e[n][i], n, N);
		if(ret > N/2) ok[n] = false;
		sub[n] += ret;
	}
	if(N-sub[n] > N/2) ok[n] = false;
	return sub[n];
}

int ans = 200000;
void DivideAndConquer(vector<int> v){
	unordered_map<int, vector<int>> g;
	rep(i, v.size()){
		g[C[v[i]]].pb(v[i]);
	}

	dfs(v[0], -1, v.size());
	int m = -1;
	rep(i, v.size()){
		if(ok[v[i]]) m = v[i];
	}
	dfs(m, -1, v.size());

	unordered_set<int> group;
	queue<int> que;
	unordered_set<int> vis;
	if(g[C[m]].size() != G[C[m]]) goto NEXT;
	rep(i, g[C[m]].size()) que.push(g[C[m]][i]);
	group.insert(C[m]);
	vis.insert(m);
	while(!que.empty()){
		int a = que.front();
		que.pop();
		while(m != a){
			a = par[a];
			if(vis.count(a)) break;
			vis.insert(a);
			if(group.count(C[a])) continue;
			if(g[C[a]].size() != G[C[a]]) goto NEXT;
			rep(i, g[C[a]].size()){
				que.push(g[C[a]][i]);
				vis.insert(g[C[a]][i]);
			}
			group.insert(C[a]);
		}
	}

	ans = min(ans, (int)group.size()-1);

NEXT:

	used[m] = true;
	unordered_set<int> st;
	st.insert(m);
	vector<int> nxt;
	rep(i, v.size()){
		if(st.count(v[i])) continue;
		nxt.clear();
		queue<int> q;
		q.push(v[i]);
		while(!q.empty()){
			int n = q.front();
			q.pop();
			nxt.pb(n);
			st.insert(n);
			rep(j, e[n].size()){
				if(st.count(e[n][j]) || used[e[n][j]]) continue;
				q.push(e[n][j]);
			}
		}
		DivideAndConquer(nxt);
	}
}

signed main(){
	scanf("%d%d", &N, &K);
	rep(i, N-1){
		int A, B; scanf("%d%d", &A, &B); A--; B--;
		e[A].pb(B);
		e[B].pb(A);
	}
	rep(i, N){
		scanf("%d", C+i); C[i]--;
		G[C[i]]++;
	}

	vector<int> v;
	rep(i, N) v.pb(i);
	DivideAndConquer(v);
	printf("%d\n", ans);
}

Compilation message (stderr)

capital_city.cpp: In function 'int dfs(int, int, int)':
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:35:2: note: in expansion of macro 'rep'
   35 |  rep(i, e[n].size()){
      |  ^~~
capital_city.cpp: In function 'void DivideAndConquer(std::vector<int>)':
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:49:2: note: in expansion of macro 'rep'
   49 |  rep(i, v.size()){
      |  ^~~
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:55:2: note: in expansion of macro 'rep'
   55 |  rep(i, v.size()){
      |  ^~~
capital_city.cpp:63:20: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   63 |  if(g[C[m]].size() != G[C[m]]) goto NEXT;
      |     ~~~~~~~~~~~~~~~^~~~~~~~~~
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:64:2: note: in expansion of macro 'rep'
   64 |  rep(i, g[C[m]].size()) que.push(g[C[m]][i]);
      |  ^~~
capital_city.cpp:75:22: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   75 |    if(g[C[a]].size() != G[C[a]]) goto NEXT;
      |       ~~~~~~~~~~~~~~~^~~~~~~~~~
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:76:4: note: in expansion of macro 'rep'
   76 |    rep(i, g[C[a]].size()){
      |    ^~~
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:92:2: note: in expansion of macro 'rep'
   92 |  rep(i, v.size()){
      |  ^~~
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'j' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:102:4: note: in expansion of macro 'rep'
  102 |    rep(j, e[n].size()){
      |    ^~~
capital_city.cpp: In function 'int main()':
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:113:2: note: in expansion of macro 'rep'
  113 |  rep(i, N-1){
      |  ^~~
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:118:2: note: in expansion of macro 'rep'
  118 |  rep(i, N){
      |  ^~~
capital_city.cpp:15:27: warning: unnecessary parentheses in declaration of 'i' [-Wparentheses]
   15 | #define rep(i, n) for(int (i)=0; (i)<(int)(n); (i)++)
      |                           ^
capital_city.cpp:124:2: note: in expansion of macro 'rep'
  124 |  rep(i, N) v.pb(i);
      |  ^~~
capital_city.cpp:112:7: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  112 |  scanf("%d%d", &N, &K);
      |  ~~~~~^~~~~~~~~~~~~~~~
capital_city.cpp:114:18: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  114 |   int A, B; scanf("%d%d", &A, &B); A--; B--;
      |             ~~~~~^~~~~~~~~~~~~~~~
capital_city.cpp:119:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  119 |   scanf("%d", C+i); C[i]--;
      |   ~~~~~^~~~~~~~~~~

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