Submission #170874

#	Time	Username	Problem	Language	Result	Execution time	Memory
170874		nobik	Split the Attractions (IOI19_split)	C++14	46 / 100	258 ms	24040 KiB

split

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

using namespace std;

struct Fau {
	int n;

	vector<int> parent;
	vector<vector<int>> bucket;

	Fau(int n): n(n), parent(n), bucket(n) {
		for (int i = 0; i < n; ++i) {
			bucket[i].push_back(i);
			parent[i] = i;
		}
	}

	int Root(int a) {
		return parent[a] == a ? a : Root(parent[a]);
	}

	int Union(int a, int b) {
		a = Root(a);
		b = Root(b);
		if (a == b)
			return a;
		
		if (bucket[a].size() > bucket[b].size())
			swap(a, b);
		
		parent[a] = b;
		bucket[b].insert(bucket[b].end(), bucket[a].begin(), bucket[a].end());
		bucket[a].clear();
		bucket[a].shrink_to_fit();

		return b;
	}

	int Size(int v) {
		return (int) bucket[v].size();
	}
};

struct Solver {
	int n, m;

	vector<vector<pair<int, int>>> graph;
	vector<int> tree_edges;
	vector<int> used, sizes;

	vector<pair<int, int>> component_sizes;

	Solver(int n, const vector<pair<int, int>>& edges, int a, int b, int c)
			: n(n), m(edges.size()), graph(n), tree_edges(m), used(n), sizes(n) {
		component_sizes.emplace_back(a, 1);
		component_sizes.emplace_back(b, 2);
		component_sizes.emplace_back(c, 3);
		sort(component_sizes.begin(), component_sizes.end());

		for (int i = 0; i < m; ++i) {
			int x = edges[i].first, y = edges[i].second;
			graph[x].emplace_back(y, i);
			graph[y].emplace_back(x, i);
		}
	}

	void FindSpanningTree(int v) {
		used[v] = 1;
		sizes[v] = 1;
		for (const auto& edge: graph[v]) {
			int to = edge.first, id = edge.second;
			if (used[to])
				continue;
			
			tree_edges[id] = 1;
			FindSpanningTree(to);
			sizes[v] += sizes[to];
		}
	}

	int FindCentroid(int v, int parent) {
		for (const auto& edge: graph[v]) {
			int to = edge.first, id = edge.second;
			if (to == parent || !tree_edges[id])
				continue;
			
			if (2 * sizes[to] > n)
				return FindCentroid(to, v);
		}

		return v;
	}

	void UniteSubtrees(int centroid, Fau& fau) {
		for (int i = 0; i < n; ++i) {
			if (i == centroid)
				continue;
			
			for (const auto& edge: graph[i]) {
				int to = edge.first, id = edge.second;
				if (to == centroid || !tree_edges[id])
					continue;
				
				fau.Union(i, to);
			}
		}
	}

	vector<int> Solve() {
		FindSpanningTree(0);
		int centroid = FindCentroid(0, -1);

		Fau fau(n);
		UniteSubtrees(centroid, fau);

		
		int max_subtree_vertex = -1, max_subtree_size = -1;
		for (int i = 0; i < n; ++i) {
			if (fau.Root(i) == i && fau.Size(i) > max_subtree_size) {
				max_subtree_size = fau.Size(i);
				max_subtree_vertex = i;
			}
		}

		for (int i = 0; i < n; ++i) {
			if (i == centroid)
				continue;
			
			for (const auto& edge: graph[i]) {
				if (max_subtree_size >= component_sizes[0].first)
					break;
				
				int to = edge.first;
				if (to == centroid)
					continue;
				
				int vertex = fau.Union(i, to);
				if (fau.Size(vertex) > max_subtree_size) {
					max_subtree_size = fau.Size(vertex);
					max_subtree_vertex = vertex;
				}
			}
		}

		// Solution was not found!
		if (max_subtree_size < component_sizes[0].first) {
			return vector<int>(n);
		}

		vector<int> result(n, component_sizes.back().second);
		Color(max_subtree_vertex, fau.bucket[max_subtree_vertex], component_sizes[0].first, component_sizes[0].second, result);
		Color(centroid, Complement(fau.bucket[max_subtree_vertex]), component_sizes[1].first, component_sizes[1].second, result);
		return result;
	}

	void Color(int v, const vector<int>& allowed_vertices, int size, int color, vector<int>& colors) {
		vector<int> used(n, 1);
		for (int e: allowed_vertices)
			used[e] = 0;

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

		while (!q.empty() && size > 0) {
			v = q.front();
			q.pop();
			colors[v] = color;
			--size;

			for (const auto& edge: graph[v]) {
				int to = edge.first;
				if (!used[to]) {
					q.push(to);
					used[to] = 1;
				}
			}
		}
	}

	vector<int> Complement(const vector<int>& a) {
		vector<int> b(n), result;
		for (int e: a) b[e] = 1;
		for (int i = 0; i < n; ++i) {
			if (!b[i])
				result.push_back(i);
		}

		return result;
	}
};

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

	Solver solver(n, edges, a, b, c);
	return solver.Solve();
}

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