Submission #1227277

#	Time	Username	Problem	Language	Result	Execution time	Memory
1227277		colossal_pepe	Split the Attractions (IOI19_split)	C++20	0 / 100	1 ms	328 KiB

split

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

const pair<int, int> null_pair = make_pair(-1, -1);

int n, root;
vector<vector<int>> t;
vector<int> abc;
vector<int> subt_n;

vector<int> split(int x, int y) {
	if (x == -1) return vector<int>(n, -1);
	vector<int> ret(n, 2);
	int c = 0;
	auto dfs = [&](const auto &self, int u, int par) -> void {
		if (abc[c]) ret[u] = c, abc[c]--;
		for (int v : t[u]) {
			if (v == par) continue;
			self(self, v, u);
		}
	};
	dfs(dfs, x, y);
	c = 1;
	dfs(dfs, y, x);
	return ret;
}

vector<int> findSplit() {
	subt_n.resize(n);
	auto dfs1 = [&](const auto &self, int u, int par) -> void {
		subt_n[u] = 1;
		for (int v : t[u]) {
			if (v == par) continue;
			self(self, v, u);
			subt_n[u] += subt_n[v];
		}
	};
	dfs1(dfs1, root, root);
	auto dfs2 = [&](const auto &self, int u, int par) -> pair<int, int> {
		if (abc[0] <= subt_n[u] and abc[1] <= n - subt_n[u]) return make_pair(u, par);
		if (abc[1] <= subt_n[u] and abc[1] <= n - subt_n[u]) return make_pair(par, u);
		for (int v : t[u]) {
			if (v == par) continue;
			auto res = self(self, v, u);
			if (res != null_pair) return res;
		}
		return null_pair;
	};
	auto [x, y] = dfs2(dfs2, root, root);
	return split(x, y);
}

vector<int> find_split(int N, int A, int B, int C, vector<int> U, vector<int> V) {
	if (U.size() != N - 1) return vector<int>(N, 0);
	n = N;
	abc = vector<int>{A, B, C};
	sort(abc.begin(), abc.end());
	t.resize(n);
	for (int i = 0; i < n - 1; i++) {
		t[U[i]].push_back(V[i]);
		t[V[i]].push_back(U[i]);
	}
	for (int i = 0; i < n; i++) {
		if (t[i].size() == 1) root = i;
	}
	vector<int> split = findSplit();
	for (int &x : split) x++;
	return split;
}

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