Submission #1241592

#TimeUsernameProblemLanguageResultExecution timeMemory
1241592The_SamuraiSplit the Attractions (IOI19_split)C++20
40 / 100
64 ms35772 KiB
#include "split.h"
#include "bits/stdc++.h"
using namespace std;
#define ff first
#define ss second

mt19937 rng(3124124);
int rand(int l, int r) {
	int x = rng();
	return abs(x) % (r - l + 1) + l;
}

vector<int> find_split(int n, int a, int b, int c, vector<int> U, vector<int> V) {
	vector<vector<int>> g(n);
	for (int i = 0; i < U.size(); i++) g[U[i]].emplace_back(V[i]);
	for (int i = 0; i < U.size(); i++) g[V[i]].emplace_back(U[i]);
	vector<int> ans(n);
	vector<pair<int, int>> shit;
	shit.emplace_back(a, 1);
	shit.emplace_back(b, 2);
	shit.emplace_back(c, 3);
	sort(shit.begin(), shit.end());
	vector<int> sub(n, 1);
	vector<bool> vis(n), vis2(n);

	auto collect = [&](auto &collect, int u, vector<int> &vec) -> void {
		vec.emplace_back(u);
		vis2[u] = true;
		for (int v: g[u]) {
			if (vis2[v]) continue;
			collect(collect, v, vec);
		}
	};

	auto dfs = [&](auto &dfs, int u, int par) -> bool {
		vis[u] = true;
		for (int v: g[u]) {
			if (vis[v]) continue;
			if (dfs(dfs, v, u)) return true;
			sub[u] += sub[v];
		}
		if (u == 0) return false;
		if (sub[u] >= shit[1].ff and n - sub[u] >= shit[0].ff) swap(shit[0], shit[1]);
		if (sub[u] < shit[0].ff or n - sub[u] < shit[1].ff) return false;
		// cout << u << ' ' << par << ' ' << sub[u] << endl;
		ans = vector<int>(n, shit[2].ss);
		vector<int> vec;
		vis2[par] = true;
		collect(collect, u, vec);
		vis2 = vector<bool>(n);
		for (int i = 0; i < shit[0].ff; i++) {
			ans[vec[i]] = shit[0].ss;
			vis2[vec[i]] = true;
		}
		vec.clear();
		collect(collect, par, vec);
		for (int i = 0; i < shit[1].ff; i++) ans[vec[i]] = shit[1].ss;
		return true;
	};
	dfs(dfs, 0, -1);
	
	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...