제출 #1235410

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1235410		countless	Split the Attractions (IOI19_split)	C++20	100 / 100	86 ms	21936 KiB

split

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

typedef long long ll;
typedef long double ld;

#define sp <<" "<<
#define endl "\n"

// sorry i hate globals
// but i will troll you with this

#define try kkwazzawazzakkwaquikkwalaquaza168zzabolazza
#define catch uvuvwevwevweonyetenyevweugwemubwemosas

vector<int> find_split(int N, int A, int B, int C, vector<int> U, vector<int> V) {
	int M = U.size();
	vector<int> T = {A, B, C};
	vector<int> o(3);
	iota(o.begin(), o.end(), 0);
	sort(o.begin(), o.end(), [&](int i, int j) {
		return T[i] < T[j];
	});

	vector<vector<int>> graph(N);
	for (int i = 0; i < M; i++) {
		graph[U[i]].push_back(V[i]);
		graph[V[i]].push_back(U[i]);
	}

	// find spanning tree
	vector<bool> vis(N);
	vector<vector<int>> tree(N);

	auto try = [&](auto &&try, int u) -> void {
		vis[u] = true;

		for (auto &v : graph[u]) {
			if (!vis[v]) {
				tree[u].push_back(v);
				tree[v].push_back(u);
				try(try, v);
			}
		}
	};

	try(try, 0);

	// find centroid of tree
	// also wow it discreetly became DSU ?!
	bool found = false;
	int cent = 0;
	vector<int> sz(N, 1), par(N);
	iota(par.begin(), par.end(), 0);
	auto catch = [&](auto catch, int u, int p, int h) -> void {
		par[u] = h;

		for (auto &v : tree[u]) {
			if (v != p) {
				if (u == cent) h = v;
				catch(catch, v, u, h);
				sz[u] += sz[v];
			}
		}

		if (!found and sz[u] * 2 >= N) {
			found = true;
			cent = u;
		}
	};

	catch(catch, cent, -1, cent);

	// find subtree sizes
	unordered_set<int> leads;
	sz.assign(N, 1);

	// sneaky golem in the pocket
	// i'm currently listening to peltorator explain FFT
	// will this be useful to me? probably not in the near future
	// am i entertained? yes, i subscribed (not paid to say this i swear)
	// (also wow did you see that? i reused the FUNC.)
	catch(catch, cent, -1, cent);

	// casework:
	// if we have a subtree of sz >= a (recall subtree size is now maxed by n/2)
	// then we can assign all of this to a, and the rest to b
	
	// if we don't then all sz < a, keep connecting via the original graph
	// until we get a sz >= a, then assign the rest to b
	int oth = -1;
	for (auto &v : tree[cent]) {
		if (sz[v] >= T[o[0]]) {
			oth = v;
			break;
		}
	}

	auto find = [&](auto &&find, int u) -> int {
		if (u == par[u]) return u;
		return par[u] = find(find, par[u]);
	};

	auto unite = [&](int u, int v) -> int {
		u = find(find, u), v = find(find, v);

		if (u == v) return u;
		if (sz[u] < sz[v]) swap(u, v);
		par[v] = u;
		sz[u] += sz[v];
		return u;
	};

	if (oth == -1) {
		for (int i = 0; i < M; i++) {
			int u = find(find, U[i]);
			int v = find(find, V[i]);

			if (u == cent or v == cent) continue;

			int w = unite(u, v);
			if (sz[w] >= T[o[0]]) {
				oth = w;
				break;
			}
		}
	}

	vector<int> res(N);
	auto trav = [&](int st, int pr, int color, int cap) -> void {
		queue<int> q;
		q.emplace(st);
		vis[st] = true;
		res[st] = color;
		cap--;

		if (cap == 0) return;

		while (!q.empty()) {
			auto u = q.front(); q.pop();

			for (auto &v : graph[u]) {
				int w = find(find, v);
				bool ok = (pr == -1 ? true : pr == w);
				if (!vis[v] and ok and cap) {
					vis[v] = true;
					res[v] = color;
					q.emplace(v);
					cap--;

					if (cap == 0) return;
				}
			}
		}
	};

	if (oth != -1) {
		vis.assign(N, false);
		res.assign(N, o[2]+1);

		trav(oth, oth, o[0]+1, T[o[0]]);
		trav(cent, -1, o[1]+1, T[o[1]]);
	}

	return res;
}

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