Submission #198060

#	Time	Username	Problem	Language	Result	Execution time	Memory
198060		Minnakhmetov	Split the Attractions (IOI19_split)	C++14	40 / 100	918 ms	1048580 KiB

split

#include "split.h"
#include <bits/stdc++.h>
using namespace std;

const int N = 2e5 + 5;
vector<int> g[N];
bool used[N];
int sz[N];

vector<int> solve_for_line_or_cyc(int n, int a, int b, int c) {
	vector<int> res(n, 0);

	int x = 0;
	while (x < n && g[x].size() > 1)
		x++;

	if (x == n)
		x = 0;

	vector<int> v;

	while (v.size() < n) {
		used[x] = 1;
		v.push_back(x);

		for (int to : g[x]) {
			if (!used[to]) {
				x = to;
				break;
			}
		}
	}

	for (int i = 0; i < a; i++) {
		res[v[i]] = 1;
	}

	for (int i = a; i < a + b; i++) {
		res[v[i]] = 2;
	}

	for (int i = a + b; i < n; i++) {
		res[v[i]] = 3;
	}

	return res;
}

vector<int> solve_for_two_sets(int n, int a, int b, int c) {
	vector<int> res(n, 0);

	int sz_b = 0;

	queue<int> q;
	q.push(0);
	used[0] = 1;

	while (!q.empty()) {
		int node = q.front();
		q.pop();
		sz_b++;
		res[node] = 2;

		if (sz_b == b)
			break;

		for (int to : g[node]) {
			if (!used[to]) {
				used[to] = 1;
				q.push(to);
			}
		}
	}

	int x = 0;
	while (res[x])
		x++;
	res[x] = 1;

	for (int i = 0; i < n; i++) {
		if (res[i] == 0)
			res[i] = 3;
	}

	return res;
}

pair<int, int> dfs(int node, int anc, int sz_a, int sz_b, int n) {
	sz[node] = 1;
	for (int to : g[node]) {
		if (to != anc) {
			auto p = dfs(to, node, sz_a, sz_b, n);
			if (p.first != -1)
				return p;
			if (sz[to] >= sz_a && n - sz[to] >= sz_b)
				return {node, to};
			sz[node] += sz[to];
		}
	}
	return {-1, -1};
}

void walk(int node, int anc, vector<int> &res, int col, int &ct) {
	if (ct == 0)
		return;
	ct--;
	res[node] = col;
	for (int to : g[node]) {
		if (to != anc) {
			walk(to, node, res, col, ct);
		}
	}
}

vector<int> solve_for_tree(int n, int a, int b, int c) {
	pair<int, int> w[3] = {{a, 1}, {b, 2}, {c, 3}};
	sort(w, w + 3);

	auto p = dfs(0, -1, w[0].first, w[1].first, n);

	if (p.first == -1) {
		swap(w[0], w[1]);
		p = dfs(0, -1, w[0].first, w[1].first, n);
	}

	vector<int> res(n, 0);

	if (p.first != -1) {
		walk(p.second, p.first, res, w[0].second, w[0].first);
		walk(p.first, p.second, res, w[1].second, w[1].first);

		for (int &x : res) {
			if (x == 0)
				x = w[2].second;
		}
	}

	return res;
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	int m = p.size();
	for (int i = 0; i < m; i++) {
		g[p[i]].push_back(q[i]);
		g[q[i]].push_back(p[i]);
	}

	bool line = true;

	for (int i = 0; i < n; i++) {
		if (g[i].size() > 2) {
			line = false;
			break;
		}
	}

	vector<int> res(n, 0);

	if (line) {
		res = solve_for_line_or_cyc(n, a, b, c);
	}
	else if (a == 1) {
		res = solve_for_two_sets(n, a, b, c);
	}
	else {
		res = solve_for_tree(n, a, b, c);
	}

	return res;
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> solve_for_line_or_cyc(int, int, int, int)':
split.cpp:22:18: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  while (v.size() < n) {
         ~~~~~~~~~^~~

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