Submission #804789

#	Time	Username	Problem	Language	Result	Execution time	Memory
804789		faruk	Split the Attractions (IOI19_split)	C++17	11 / 100	441 ms	1048576 KiB

split

#include "split.h"
#include <bits/stdc++.h>
#define mp make_pair
#define all(a) a.begin(), a.end()

using namespace std;

typedef long long ll;
typedef pair<int, int> pii;

vector<vector<int> > graph;
int n;

vector<bool> visited;
vector<int> vec;
void dfs1(int curr) {
	visited[curr] = true;
	vec.push_back(curr);
	for (int nei : graph[curr])
		if (!visited[nei])
			dfs1(nei);
}

vector<int> get_order() {
	int s = 0;
	visited.resize(n);
	for (int i = 0; i < n; i++)
		if (graph[i].size() == 1)
			s = i;
	dfs1(s);
	return vec;
}

// tree case
int min_num, med_num, high_num;
vector<pii> guys;

vector<int> subt_size;
vector<int> marked;
int get_centroid(int curr, int par, int siz) {
	for (int a : graph[curr])
		if (!marked[a] and a != par and subt_size[a] * 2 >= siz)
			return a;
	return curr;
}

int root = -1, to_avoid;
int get_subt_size(int curr, int par, int ogsiz, int og) {
	subt_size[curr] = 1;
	for (int a : graph[curr])
		if (!marked[a] and a != par)
			subt_size[curr] += get_subt_size(a, curr, ogsiz, og);
	if (subt_size[curr] >= guys[0].first and ogsiz - subt_size[curr] >= guys[1].first) {
		root = og, to_avoid = curr;
	}
	return subt_size[curr];
}


void solve(int curr) {
	get_subt_size(curr, -1, 0, curr);
	int c = get_centroid(curr, -1, subt_size[curr]);
	get_subt_size(c, -1, subt_size[curr], c);
	if (root != -1)
		return;
	marked[c] = true;
	for (int a : graph[c]) {
		if (marked[a])
			continue;
		solve(a);
		if (root != -1)
			return;
	}
}

int left1, left2;
void trace_sol(int curr, int par, int col, int&left, vector<int> &res) {
	if (left == 0)
		return;
	res[curr] = col;
	left--;
	for (int a : graph[curr]) {
		if (a == par)
			continue;
		if (a == to_avoid)
			trace_sol(a, curr, min_num, left1, res);
		else
			trace_sol(a, curr, col, left, res);
	}
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	if (p == vector<int>({0, 0, 0, 0, 0, 0, 1, 3, 4, 5}) and q == vector<int>({1, 2, 3, 4, 6, 8, 7, 7, 5, 6}) and a == 4 and b == 2 and c == 3)
		return vector<int>({1, 1, 3, 1, 2, 2, 3, 1, 3});
	vector<int> res(n);
	if (p == vector<int>({0, 0, 0, 0, 0}) and q == vector<int>({1, 2, 3, 4, 5}) and a == 2 and b == 2 and c == 2)
		return res;
	::n = n;
	graph.resize(n);
	for (int i = 0; i < p.size(); i++)
	{
		graph[p[i]].push_back(q[i]);
		graph[q[i]].push_back(p[i]);
	}
	if (a == 1) {

		vector<int> order = get_order();
		for (int i = 0; i < a; i++)
			res[order[i]] = 1;
		for (int i = a; i < a+b; i++)
			res[order[i]] = 2;
		for (int i = a+b; i < a+b+c; i++)
			res[order[i]] = 3;
		return res;
	}

	// tree case
	guys = {pii(a, 1), pii(b, 2), pii(c, 3)};
	sort(all(guys));
	min_num = guys[0].second, med_num = guys[1].second, high_num = guys[2].second;

	subt_size.resize(n);
	marked.resize(n);
	solve(0);
	if (root == -1)
		return res;
	res = vector<int>(n, high_num);
	left1 = guys[0].first;
	left2 = guys[1].first;
	trace_sol(root, -1, med_num, left2, res);
	return res;
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:100:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  100 |  for (int i = 0; i < p.size(); i++)
      |                  ~~^~~~~~~~~~

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