Submission #1290808

#TimeUsernameProblemLanguageResultExecution timeMemory
1290808kahoulSplit the Attractions (IOI19_split)C++20
0 / 100
606 ms1114112 KiB
	#include <algorithm>
	#include "split.h"

	const int maxn = 2e5 + 10;
	std::vector<int> rel[maxn];
	int deg[maxn];
	int pai[maxn];

	std::vector<int> res(maxn, 3);
	std::vector<bool> used(maxn, 0);

	int a, b, c;
	int needed = 0;
	int amt = 0;
	int amt_resto = 0;
	int sz[maxn];

	void dfs1 (int u, int p) {
		sz[u] = 1;
		pai[u] = p;

		for (auto v : rel[u]) {
			if (v == p) continue;
			dfs1(v, u);
			sz[u] += sz[v];
		}
	}

	void dfs (int u, int p, int color) {
		if (amt == needed) return;

		used[u] = 1;
		res[u] = color;
		amt++;

		for (auto v : rel[u]) {
			if (v == pai[u]) continue;
			if (used[v]) continue;
			if (amt == needed) return;
			dfs(v, u, color);
		}
	}

	void dfs_resto (int u, int p, int color) {
		used[u] = 1;
		res[u] = color;

		for (auto v : rel[u]) {
			if (used[v]) continue;
			dfs_resto(v, u, color);
		}
	}

	std::vector<int> find_split(int n, int A, int B, int C, std::vector<int> p, std::vector<int> q) {
		a = A;
		b = B;
		c = C;

		int cor_real[4];

		for (int i = 0; i < p.size(); i++) {
			int u = p[i];
			int v = q[i];
			rel[u].push_back(v);
			rel[v].push_back(u);
			deg[u]++;
			deg[v]++;
		}

		std::vector<std::pair<int, int>> colors = {{a, 1}, {b, 2}, {c, 3}};
		std::sort(colors.begin(), colors.end());
		cor_real[3] = colors[2].second;

		dfs1(0, 0);

		int pika = -1;
		int cor = -1;
		int cor_dois = -1;

		for (int i = 0; i < n; i++) {
			//std::cerr << sz[i] << ' ';
		}
		//std::cerr << '\n';

		for (int i = 0; i < n; i++) {
			int leftover = sz[0] - sz[i];
			if (colors[0].first <= sz[i] && colors[1].first <= leftover) {
				pika = i;
				cor = 1;
				cor_dois = 2;
				amt = colors[0].first;
				amt_resto = colors[1].first;
				cor_real[1] = colors[0].second;
				cor_real[2] = colors[1].second;
			}
			if (colors[1].first <= sz[i] && colors[0].first <= leftover) {
				pika = i;
				cor = 2;
				cor_dois = 1;
				amt = colors[1].first;
				amt_resto = colors[0].first;
				cor_real[1] = colors[1].second;
				cor_real[2] = colors[0].second;
			}
		}

		if (pika == -1) {
			std::vector<int> ans;

			for (int i = 0; i < n; i++) {
				ans.push_back(0);
			}

			return ans;
		}

		dfs(pika, pika, cor);

		needed = amt_resto;
		amt = 0;

		dfs(0, 0, cor_dois);

		std::vector<int> ans;

		for (int i = 0; i < n; i++) {
			ans.push_back(cor_real[res[i]]);
		}

		return ans;
	}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...