Submission #298669

#	Time	Username	Problem	Language	Result	Execution time	Memory
298669		keko37	Split the Attractions (IOI19_split)	C++14	100 / 100	250 ms	19428 KiB

split

#include<bits/stdc++.h>
#include "split.h"

using namespace std;

typedef vector <int> vi;
typedef pair <int, int> pi;

const int MAXN = 200005;

int n, m, root, cnt;
pi col[3];
int par[MAXN], siz[MAXN], sol[MAXN];
vector <int> v[MAXN];
vector <pi> e;
set <pi> st;

int nadi (int x) {
	if (x == par[x]) return x;
	return par[x] = nadi(par[x]);
}

bool spoji (int a, int b) {
	a = nadi(a); b = nadi(b);
	if (a == b) return 0;
	par[a] = b;
	siz[b] += siz[a];
	return 1;
}

void dfs (int x, int rod) {
	siz[x] = 1;
	for (auto sus : v[x]) {
		if (sus == rod) continue;
		dfs(sus, x);
		siz[x] += siz[sus];
	}
}

int nadi_centroid (int x, int rod) {
	int mx = 0, ind = 0;
	for (auto sus : v[x]) {
		if (sus == rod) continue;
		if (siz[sus] > mx) {
			mx = siz[sus];
			ind = sus;
		}
	}
	if (mx * 2 <= n) return x;
	return nadi_centroid(ind, x);
}

bool check () {
	if (st.empty()) return 0;
	auto it = st.end(); it--;
	return (it -> first) >= col[0].first;
}

void oboji (int x) {
	if (cnt == col[1].first) return;
	sol[x] = col[1].second;
	cnt++;
	for (auto sus : v[x]) {
		if (cnt == col[1].first) return;
		if (sol[sus] == 0) oboji(sus);
	}
}

int p;

void small (int x) {
	if (cnt == col[0].first) return;
	sol[x] = col[0].second;
	cnt++;
	for (auto sus : v[x]) {
		if (cnt == col[0].first) return;
		if (sol[sus] == 0 && nadi(sus) == p) small(sus);
	}
}

vi solve () {
	auto it = st.end(); it--;
	p = it -> second;
	cnt = 0; small(p);
	cnt = 0; oboji(root);
	vi res(n);
	for (int i = 1; i <= n; i++) {
		if (sol[i] == 0) sol[i] = col[2].second;
		res[i - 1] = sol[i];
	}
	return res;
}

vi find_split (int N, int A, int B, int C, vi P, vi Q) {
	n = N; m = P.size();
	col[0] = {A, 1}; col[1] = {B, 2}; col[2] = {C, 3};
	sort(col, col + 3);
	for (int i = 1; i <= n; i++) {
		par[i] = i;
	}
	for (int i = 0; i < m; i++) {
		P[i]++; Q[i]++;
		if (spoji(P[i], Q[i])) {
			v[P[i]].push_back(Q[i]);
			v[Q[i]].push_back(P[i]);
		} else {
			e.push_back({P[i], Q[i]});
		}
	}
	dfs(1, 0);
	root = nadi_centroid(1, 0);
	for (int i = 1; i <= n; i++) {
		par[i] = i;
		siz[i] = 1;
	}
	for (int i = 1; i <= n; i++) {
		for (auto sus : v[i]) {
			if (i != root && sus != root && i < sus) spoji(i, sus);
		}
	}
	for (auto sus : v[root]) {
		st.insert({siz[nadi(sus)], nadi(sus)});
	}
	if (check()) return solve();
	for (auto edge : e) {
		int a = edge.first, b = edge.second;
		int pa = nadi(a), pb = nadi(b);
		if (a == root || b == root || pa == pb) continue;
		st.erase({siz[pa], pa});
		st.erase({siz[pb], pb});
		par[pa] = pb;
		siz[pb] += siz[pa];
		st.insert({siz[pb], pb});
		v[a].push_back(b);
		v[b].push_back(a);
		if (check()) return solve();
	}
	vi res(n);
	return res;
}

• Subtask 1: 7 / 7
• Subtask 2: 11 / 11
• Subtask 3: 22 / 22
• Subtask 4: 24 / 24
• Subtask 5: 36 / 36