Submission #294828

#	Time	Username	Problem	Language	Result	Execution time	Memory
294828		mode149256	Split the Attractions (IOI19_split)	C++14	40 / 100	149 ms	15992 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;
vii edges(MX);
vi col(3);
vi kiek(3);
vi sz(MX, 0);
vi res;
vb been(MX, false);

bool colored = false;

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

	for (auto u : edges[x]) {
		if (been[u]) continue;
		makeSz(u);
		sz[x] += sz[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)
			color(u, x, c);
}

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 dfs(int x, int p) {
	// rooted at x
	for (auto u : edges[x]) {
		if (sz[u] >= kiek[0] and N - sz[u] >= kiek[1]) {
			colored = true;
			color(u, x, 0);
			color(x, u, 1);
			return;
		} else if (sz[u] >= kiek[1] and N - sz[u] >= kiek[0]) {
			colored = true;
			color(u, x, 1);
			color(x, u, 0);
			return;
		}
	}

	for (auto u : edges[x]) {
		if (u == p) 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 < 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) {
		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