Submission #143748

#	Time	Username	Problem	Language	Result	Execution time	Memory
143748		joseacaz	Split the Attractions (IOI19_split)	C++17	40 / 100	169 ms	16340 KiB

split

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

using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;

const int MAXN = 100005;

int N, M, sz[MAXN], parent[MAXN], nodes, nodeA, vis[MAXN], vis2[MAXN], centroid;
vi Tree[MAXN], Graph[MAXN], ans, nodes2;
vector < pii > part ( 3, {0, 0} );
map<pii, int> MP;

int pre ( int node, int p = -1 )
{
	sz[node] = 1;
	for ( auto i : Tree[node] )
		if ( i != p )
			sz[node] += pre ( i, node );
	return sz[node];
}

int find_centroid ( int node, int p = -1 )
{
	for ( auto i : Tree[node] )
		if ( i != p && 2 * sz[i] > sz[0] )
			return find_centroid ( i, node );
	return node;
}

int look ( int node )
{
	if ( node == parent[node] )
		return node;
	return parent[node] = look ( parent[node] );
}

void join ( int nodeA, int nodeB )
{
	parent[look(nodeB)] = look(nodeA);
}

void dfs ( int root, int node, int p = -1 )
{
	join ( root, node );
	for ( auto i : Tree[node] )
		if ( i != p )
			dfs ( root, i, node );
}

void paint ( int node, int color )
{
	if ( nodes )
		ans[node] = color, nodes--;
	else
		ans[node] = part[2].second;
	
	vis[node] = 1;
	for ( auto i : Tree[node] )
		if ( !vis[i] )
			paint ( i, color );
}

void possible ( int node )
{
	nodes -= sz[node];
	vis2[node] = 1;
	for ( auto i : Graph[node] )
		if ( !vis2[i] )
			possible ( i );
}

void paint2 ( int node, int color )
{
	if ( nodes > 0 )
	{
		paint ( node, color );
		vis[node] = 1;
	}
	
	vis2[node] = 1;
	for ( auto i : Graph[node] )
		if ( !vis2[i] )
			paint2 ( i, color );
}

vi find_split ( int _n, int a, int b, int c, vi p, vi q )
{
	part[0] = {a, 1};
	part[1] = {b, 2};
	part[2] = {c, 3};
	sort ( part.begin(), part.end() );
	N = _n;
	M = p.size();

	ans.resize ( N );
	for ( int i = 0; i < N; i++ )
		parent[i] = i, ans[i] = 0;
	for ( int i = 0; i < M; i++ )
	{
		if ( look ( p[i] ) != look ( q[i] ) )
		{
			join ( p[i], q[i] );
			Tree[p[i]].push_back ( q[i] );
			Tree[q[i]].push_back ( p[i] );
		}
	}
	
	pre ( 0 );
	centroid = find_centroid ( 0 );
	vis[centroid] = 1;
	pre ( centroid );
	for ( int i = 0; i < N; i++ )
		parent[i] = i;
	
	nodeA = -1;
	for ( auto i : Tree[centroid] )
	{
		dfs ( i, i, centroid );
		if ( sz[i] >= part[0].first )
			nodeA = i;
		nodes2.push_back ( i );
	}
	
	if ( nodeA != -1 )
	{
		vis[nodeA] = 1;
		nodes = part[0].first;
		paint ( nodeA, part[0].second );
		nodes = part[1].first;
		paint ( centroid, part[1].second );
		return ans;
	}

	for ( int i = 0; i < M; i++ )
	{
		if ( p[i] == centroid or q[i] == centroid )
			continue;

		if ( look(p[i]) != look(q[i]) && !MP[{look(p[i]), look(q[i])}])
		{
			Graph[look(p[i])].push_back ( look(q[i]) );
			Graph[look(q[i])].push_back ( look(p[i]) );
			MP[{look(p[i]), look(q[i])}] = 1;
			MP[{look(q[i]), look(p[i])}] = 1;
		}
	}
	
	for ( auto i : nodes2 )
	{
		nodes = part[0].first;
		possible ( i );
		
		if ( nodes <= 0 )
		{
			for ( auto j : nodes2 )
				vis2[j] = 0;
			nodes = part[0].first;
			paint2 ( i, part[0].second );
			nodes = part[1].first;
			paint ( centroid, part[1].second );
			return ans;
		}
	}
	
	for ( auto &i : ans )
		i = 0;
	return ans;
}

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