Submission #295488

#	Time	Username	Problem	Language	Result	Execution time	Memory
295488		mode149256	Split the Attractions (IOI19_split)	C++14	40 / 100	255 ms	49656 KiB

split

#include<bits/stdc++.h>
#include "split.h"
using namespace std;
#define x first
#define y second

using ll = long long;
using vi = vector<int>;
using vl = vector<ll>;
using vb = vector<bool>;
using vii = vector<vi>;
using pi = pair<int, int>;
using vpi = vector<pi>;

const int MX = 2e5 + 100;

int N, M;
int A, B, C;
vpi visos;
vii edges(MX);
vii zemyn(MX);
unordered_set<int> inTree[MX];
vi col(3);
vi kiek(3);
vi sz(MX, 0);
vi res;
vb been(MX, false);
vi par(MX);
vi szp(MX, 1);

int findSet(int i) { return (i == par[i] ? i : (par[i] = findSet(par[i]))); }
bool isSameSet(int i, int j) { return findSet(i) == findSet(j); }
bool unionSet(int i, int j) {
	if (isSameSet(i, j)) return false;

	i = findSet(i);
	j = findSet(j);

	if (szp[i] > szp[j]) swap(i, j);

	par[i] = j;
	szp[j] += szp[i];
	return true;
}

bool colored = false;
bool nepilnas = false;

void makeSz(int x) {
	sz[x] = 1;
	been[x] = true;

	for (auto u : edges[x]) {
		if (been[u]) continue;
		inTree[x].insert(u);
		inTree[u].insert(x);
		// zemyn[x].emplace_back(u);

		makeSz(u);
		sz[x] += sz[u];
	}
}

void makeZemyn(int x) {
	been[x] = true;
	for (auto u : edges[x]) {
		if (been[u] or !inTree[x].count(u)) continue;
		zemyn[x].emplace_back(u);
		// printf("zemyn x -> u %d -> %d\n", x, u);
		makeZemyn(u);
	}
}

void color(int x, int p, int c) {
	if (kiek[c] <= 0) return;
	// printf("c = %d, x = %d, col = %d, kiek = %d\n", c, x, col[c], kiek[c]);
	res[x] = col[c];
	kiek[c]--;
	if (kiek[c] <= 0) return;

	for (auto u : edges[x])
		if (u != p and res[u] == 0)
			color(u, x, c);
}

void colorMult(int x, int c, int cent) {
	if (kiek[c] <= 0) return;
	res[x] = col[c];
	kiek[c]--;
	if (kiek[c] <= 0) return;

	for (auto u : edges[x]) {
		if (u == cent) continue;
		if (inTree[x].count(u) == 0) continue;
		if (res[u] != 0) continue;
		colorMult(u, c, cent);
	}
}

void colorZemyn(int x, int c) {
	if (kiek[c] <= 0) return;
	res[x] = col[c];
	kiek[c]--;
	if (kiek[c] <= 0) return;

	for (auto u : zemyn[x])
		colorZemyn(u, c);
}

void unite(int x) {
	for (auto u : zemyn[x]) {
		// printf("jungiu u = %d, x = %d\n", u, x);
		unionSet(x, u);
		unite(u);
	}
}

void colorAll(int x, int c) {
	been[x] = true;
	if (kiek[c] <= 0) return;
	// printf("c = %d, x = %d, col = %d, kiek = %d\n", c, x, col[c], kiek[c]);
	res[x] = col[c];
	kiek[c]--;
	if (kiek[c] <= 0) return;

	for (auto u : edges[x])
		if (!been[u])
			colorAll(u, c);
}

void solve(int x) {
	// x is centroid
	been = vb(N, false);
	makeZemyn(x);

	for (auto u : edges[x]) {
		unite(u);
		// printf("x = %d, u = %d, szu = %d, szp = %d\n", x, u, sz[u], szp[findSet(u)]);
		// assert(sz[u] == szp[findSet(u)]);

		if (sz[u] >= kiek[0]) {
#ifdef LOCAL
			printf("PIRMO x = %d, u = %d\n", x, u);
#endif
			colorZemyn(u, 0);
			color(x, u, 1);
			return;
		}
	}

	// reik jungt
	for (auto u : visos) {
		if (u.x == x or u.y == x) continue;
		if (isSameSet(u.x, u.y)) continue;

		inTree[u.x].insert(u.y);
		inTree[u.y].insert(u.x);

		unionSet(u.x, u.y);

		if (szp[findSet(u.x)] >= kiek[0]) {
#ifdef LOCAL
			printf("SUJUNGUS x = %d, u = %d %d\n", x, u.x, u.y);
#endif
			colorMult(findSet(u.x), 0, x);
			color(x, -1, 1);
			return;
		}
	}

	nepilnas = true;
}

void dfs(int x, int p) {
	int centroid = 1;
	for (auto u : edges[x]) {
		if (sz[u] > N / 2)
			centroid = 0;
	}

	if (centroid) {
		colored = true;
		solve(x);
		return;
	}

	for (auto u : edges[x]) {
		if (u == p or !inTree[x].count(u)) continue;

		int prevX = sz[x];
		int prevU = sz[u];

		sz[x] -= sz[u];
		sz[u] += sz[x];
		dfs(u, x);

		sz[x] = prevX;
		sz[u] = prevU;
		if (colored) return;
	}
}

void sortabc() {
	if (A >= B and A >= C) {
		col[2] = 1;
		kiek[2] = A;

		col[1] = 2;
		kiek[1] = B;

		col[0] = 3;
		kiek[0] = C;
	} else if (B >= A and B >= C) {
		col[2] = 2;
		kiek[2] = B;

		col[1] = 1;
		kiek[1] = A;

		col[0] = 3;
		kiek[0] = C;
	} else {
		col[2] = 3;
		kiek[2] = C;

		col[1] = 2;
		kiek[1] = B;

		col[0] = 1;
		kiek[0] = A;
	}

	if (kiek[0] > kiek[1]) {
		swap(kiek[0], kiek[1]);
		swap(col[0], col[1]);
	}
}

vi find_split(int n, int a, int b, int c, vi p, vi q) {
	res.resize(n, 0);
	N = n;
	A = a;
	B = b;
	C = c;
	M = (int)p.size();
	for (int i = 0; i < N; ++i)
		par[i] = i;
	for (int i = 0; i < M; ++i)
	{
		edges[p[i]].emplace_back(q[i]);
		edges[q[i]].emplace_back(p[i]);
	}

	sortabc();

	// if (kiek[0] == 1) {
	// 	been.resize(N, false);
	// 	colorAll(0, 1);
	// 	for (int i = 0; i < N; ++i)
	// 	{
	// 		if (!res[i]) {
	// 			if (kiek[0]) {
	// 				res[i] = col[0];
	// 				kiek[0]--;
	// 			} else {
	// 				res[i] = col[2];
	// 			}
	// 		}
	// 	}

	// 	return res;
	// }

	been.resize(N, false);
	makeSz(0);
	been.resize(N, false);
	dfs(0, -1);

	// for (int i = 0; i < N; ++i)
	// {
	// 	printf("%d ", res[i]);
	// }
	// printf("\n");
	// printf("col %d %d %d\n", col[0], col[1], col[2]);
	if (colored and !nepilnas) {
		for (int i = 0; i < N; ++i)
		{
			if (!res[i]) res[i] = col[2];
		}
	}
	return res;
}

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