Submission #421866

#	Time	Username	Problem	Language	Result	Execution time	Memory
421866		KoD	Split the Attractions (IOI19_split)	C++17	100 / 100	181 ms	16684 KiB

split

#include <bits/stdc++.h>
#include "split.h"

template <class T> using Vec = std::vector<T>;

template <class F> struct RecLambda: private F {
	explicit RecLambda(F&& f): F(std::forward<F>(f)) {}
	template <class... Args> decltype(auto) operator () (Args&&... args) const {
		return F::operator()(*this, std::forward<Args>(args)...);
	}
};

struct DSU {
	Vec<int> par;
	DSU(const int n): par(n, -1) {}
	int find(const int x) {
		return par[x] < 0 ? x : par[x] = find(par[x]);
	}
	bool merge(int x, int y) {
		x = find(x);
		y = find(y);
		if (x != y) {
			par[y] += par[x];
			par[x] = y;
			return true;
		}
		return false;
	}
};

Vec<int> find_split(int n, int a, int b, int c, Vec<int> x, Vec<int> y) {
	std::array<int, 4> specified_size = {-1, a, b, c};
	std::array<int, 3> size_order = {1, 2, 3};
	std::sort(size_order.begin(), size_order.end(), [&](const int i, const int j) {
		return specified_size[i] < specified_size[j];
	});
	const auto [small, medium, big] = size_order;
	const int m = (int) x.size();
	DSU dsu(n);
	Vec<std::pair<int, int>> tree;
	Vec<std::pair<int, int>> other;
	for (int i = 0; i < m; ++i) {
		const auto u = x[i];
		const auto v = y[i];
		if (dsu.merge(u, v)) {
			tree.emplace_back(u, v);
		} else {
			other.emplace_back(u, v);
		}
	}
	Vec<Vec<int>> graph(n);
	for (const auto [u, v]: tree) {
		graph[u].push_back(v);
		graph[v].push_back(u);
	}
	Vec<int> subtree(n);
	RecLambda([&](auto&& dfs, const int u, const int p) -> void {
		subtree[u] = 1;
		for (const auto v: graph[u]) {
			if (v != p) {
				dfs(v, u);
				subtree[u] += subtree[v];
			}
		}
	})(0, -1);
	Vec<int> ret(n);
	const auto write_tree = RecLambda([&](auto&& dfs, const int u, const int p, const int w, int& c) -> void {
		if (c == 0) {
			return;
		}
		ret[u] = w;
		c -= 1;
		for (const auto v: graph[u]) {
			if (v != p) {
				dfs(v, u, w, c);
			}
		}
	});
	bool done = false;
	for (auto& [u, v]: tree) {
		if (subtree[u] < subtree[v]) {
			std::swap(u, v);
		}
		if (subtree[v] >= specified_size[small] and n - subtree[v] >= specified_size[medium]) {
			write_tree(v, u, small, specified_size[small]);
			write_tree(u, v, medium, specified_size[medium]);
			done = true;
			break;
		}
		if (subtree[v] >= specified_size[medium] and n - subtree[v] >= specified_size[small]) {
			write_tree(u, v, small, specified_size[small]);
			write_tree(v, u, medium, specified_size[medium]);
			done = true;
			break;
		}
	}
	if (!done) {
		const auto cent = RecLambda([&](auto&& dfs, const int u, const int p) -> int {
			for (const auto v: graph[u]) {
				if (v != p and subtree[v] * 2 > n) {
					return dfs(v, u);
				}
			}
			return u;
		})(0, -1);
		DSU split(n);
		for (const auto [u, v]: tree) {
			if (u != cent and v != cent) {
				split.merge(u, v);
			}
		}
		for (const auto [u, v]: other) {
			if (u != cent and v != cent) {
				split.merge(u, v);
				const int r = split.find(u);
				if (-split.par[r] >= specified_size[small]) {
					for (const auto [u, v]: other) {
						graph[u].push_back(v);
						graph[v].push_back(u);
					}
					RecLambda([&](auto&& dfs, const int u) -> void {
						if (specified_size[small] == 0 or ret[u] != 0 or split.find(u) != r) {
							return;
						}
						ret[u] = small;	
						specified_size[small] -= 1;
						for (const auto v: graph[u]) {
							dfs(v);
						}
					})(r);
					RecLambda([&](auto&& dfs, const int u) -> void {
						if (specified_size[medium] == 0 or ret[u] != 0 or split.find(u) == r) {
							return;
						}
						ret[u] = medium;
						specified_size[medium] -= 1;
						for (const auto v: graph[u]) {
							dfs(v);
						}
					})(cent);
					break;
				}
			}
		}
	}
	if (std::any_of(ret.begin(), ret.end(), [&](const int x) { return x != 0; })) {
		for (auto& x: ret) {
			if (x == 0) {
				x = big;
			}
		}
	}
	return ret;	
}	

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