Submission #1066193

#	Time	Username	Problem	Language	Result	Execution time	Memory
1066193		Tsovak	Split the Attractions (IOI19_split)	C++17	40 / 100	534 ms	1048576 KiB

split

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

// -------------------- Typedefs -------------------- //

typedef long long ll;
typedef unsigned long long ull;
typedef double db;
typedef long double ld;

// -------------------- Defines -------------------- //

#define pr pair
#define mpr make_pair
#define ff first
#define ss second

#define mset multiset
#define mmap multimap
#define uset unordered_set
#define umap unordered_map
#define umset unordered_multiset
#define ummap unordered_multimap
#define pqueue priority_queue

#define sz(x) (int((x).size()))
#define len(x) (int((x).length()))
#define all(x) (x).begin(), (x).end()
#define clr(x) (x).clear()

#define ft front
#define bk back
#define pf push_front
#define pb push_back
#define popf pop_front
#define popb pop_back

#define lb lower_bound
#define ub upper_bound
#define bs binary_search

// -------------------- Constants -------------------- //

const int MAX = int(1e9 + 5);
const ll MAXL = ll(1e18 + 5);
const ll MOD = ll(1e9 + 7);
const ll MOD2 = ll(998244353);

// -------------------- Solution -------------------- //

const int N = 100005;
vector<int> g[N];
int sz[N], deg[N];
int ans[N];
int n, m;

int timer, tin[N], tout[N];

void dfs(int u, int p = 0)
{
	int v;

	tin[u] = ++timer;

	for (int i = 0; i < sz(g[u]); i++) {
		v = g[u][i];
		if (v == p) continue;

		dfs(v, u);
		sz[u] += sz[v];
	}
	sz[u]++;

	tout[u] = ++timer;

	return;
}

bool used[N];
int h;

void dfs2(int u)
{
	int v;

	used[u] = true;

	if (!h) return;
	ans[u] = 2;
	h--;

	for (int i = 0; i < sz(g[u]); i++) {
		v = g[u][i];
		if (used[v]) continue;

		dfs2(v);
		if (!h) return;
	}

	return;
}

bool is_anc(int u, int v)
{
	return tin[u] <= tin[v] && tout[v] <= tout[u];
}

void check(int a, int b)
{
	int i, j;
	int u, v;

	for (u = 1; u <= n; u++) if (sz[u] >= a && n - sz[u] >= b) break;

	if (u == n + 1) return;

	for (i = 1; i <= n; i++) {
		if (is_anc(u, i)) ans[i] = 1;
		else ans[i] = 2;
	}

	set<pr<int, int>> st;

	for (i = 1; i <= n; i++) {
		if (ans[i] != 1) continue;

		for (j = 0; j < sz(g[i]); j++) deg[i] += (ans[g[i][j]] == 1);

		st.insert(mpr(deg[i], i));
	}

	while (sz(st) > a) {
		u = (*st.begin()).ss;
		st.erase(st.begin());

		ans[u] = 3;

		for (j = 0; j < sz(g[u]); j++) {
			v = g[u][j];

			if (ans[v] != 1) continue;

			st.erase(st.find(mpr(deg[v], v)));
			deg[v]--;
			st.insert(mpr(deg[v], v));
		}
	}

	clr(st);

	for (i = 1; i <= n; i++) {
		if (ans[i] != 2) continue;

		deg[i] = 0;
		for (j = 0; j < sz(g[i]); j++) deg[i] += (ans[g[i][j]] == 2);

		st.insert(mpr(deg[i], i));
	}

	while (sz(st) > b) {
		u = (*st.begin()).ss;
		st.erase(st.begin());

		ans[u] = 3;

		for (j = 0; j < sz(g[u]); j++) {
			v = g[u][j];

			if (ans[v] != 2) continue;

			st.erase(st.find(mpr(deg[v], v)));
			deg[v]--;
			st.insert(mpr(deg[v], v));
		}
	}

	return;
}

vector<int> find_split(int n0, int a, int b, int c, vector<int> p0, vector<int> q0)
{
	int i, j;

	n = n0; m = sz(p0);

	for (i = 0; i < m; i++) {
		g[p0[i] + 1].pb(q0[i] + 1);
		g[q0[i] + 1].pb(p0[i] + 1);
	}

	if (a == 1) {
		for (i = 1; i <= n; i++) ans[i] = 3;

		h = b;

		dfs2(1);

		for (i = 1; i <= n; i++) {
			if (ans[i] == 3) {
				ans[i] = 1;
				break;
			}
		}

		vector<int> res(n);
		for (i = 0; i < n; i++) res[i] = ans[i + 1];

		return res;
	}

	if (m == n) {
		g[p0.bk() + 1].popb();
		g[q0.bk() + 1].popb();
	}

	dfs(1);

	check(a, b);
	if (!ans[1]) check(a, c);
	if (!ans[1]) check(b, a);
	if (!ans[1]) check(b, c);
	if (!ans[1]) check(c, a);
	if (!ans[1]) check(c, b);

	vector<int> res(n);

	for (i = 0; i < n; i++) res[i] = ans[i + 1];


	if (!ans[1]) return res;

	int cnt[4] = { 0, 0, 0, 0 }, to[4] = { 0, 1, 2, 3 }, rev[4];

	for (i = 0; i < n; i++) cnt[res[i]]++;

	do {
		if (cnt[to[1]] == a && cnt[to[2]] == b && cnt[to[3]] == c) break;
	} while (next_permutation(to + 1, to + 4));

	for (i = 1; i <= 3; i++) rev[to[i]] = i;

	for (i = 0; i < n; i++) res[i] = rev[res[i]];

	return res;
}

/*
	# # # #   # # # #   # # # #   #       #    #       #     #
	   #      #         #     #    #     #    # #      #   #
	   #      # # # #   #     #     #   #    #   #     # #
	   #            #   #     #      # #    # # # #    #   #
	   #      # # # #   # # # #       #    #       #   #     #
*/

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:185:9: warning: unused variable 'j' [-Wunused-variable]
  185 |  int i, j;
      |         ^

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