Submission #1165546

#	Time	Username	Problem	Language	Result	Execution time	Memory
1165546		hyakup	Split the Attractions (IOI19_split)	C++20	100 / 100	115 ms	19780 KiB

split

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

using namespace std;

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	vector<vector<int>> adj(n);
	vector<int> prof(n), low(n), sub(n);

	for( int i = 0; i < p.size(); i++ ){
		adj[p[i]].push_back(q[i]);
		adj[q[i]].push_back(p[i]);
	}

	function<int(int, int, int)> dfs_init = [&]( int cur, int pai, int nivel ){
		prof[cur] = low[cur] = nivel;
		sub[cur] = 1;
		for( int viz : adj[cur] ) if( viz != pai ){
			if( !prof[viz] ) sub[cur] += dfs_init( viz, cur, nivel + 1 );
			low[cur] = min( low[cur], (( prof[viz] < prof[cur] ) ? prof[viz] : low[viz] ) );
		}
		return sub[cur];
	};

	function<int(int, int)> find_centroid = [&]( int cur, int tam ){
		for( int viz : adj[cur] ) if( prof[viz] == prof[cur] + 1 && sub[viz] > tam/2 ) return find_centroid( viz, tam );
		return cur;
	};

	dfs_init( 0, 0, 1 );
	int centroid = find_centroid( 0, n );


	vector<pair<int, int>> comp = { make_pair( a, 1 ), make_pair( b, 2 ), make_pair( c, 3 ) };
	sort( comp.begin(), comp.end() );

	vector<int> resp(n, comp[2].second);

	function<void(int, int, int&, int, int)> dfs = [&]( int cur, int pai, int &qtd, int flag, int t ){
		if( qtd == 0 || resp[cur] != comp[2].second ) return;
		resp[cur] = flag; qtd--;
		for( int viz : adj[cur] ) if( t || (abs(prof[viz] - prof[cur]) == 1 && viz != pai) ) dfs( viz, cur, qtd, flag, t);
	};

	for( int viz : adj[centroid] ){
		if( prof[viz] == prof[centroid] + 1 && sub[viz] >= comp[0].first ){
			dfs( viz, centroid, comp[0].first, comp[0].second, 0 );
			dfs( centroid, viz, comp[1].first, comp[1].second, 0 );
			return resp;
		}
		if( prof[centroid] == prof[viz] + 1 && n - sub[centroid] >= comp[0].first ){
			dfs( viz, centroid, comp[0].first, comp[0].second, 0 );
			dfs( centroid, viz, comp[1].first, comp[1].second, 0 );
			return resp;
		}
	}

	int total = n - sub[centroid];
	resp[centroid] = comp[1].second;
	comp[1].first--;
	for( int viz : adj[centroid] ) if( prof[viz] == prof[centroid] + 1 ){
		if( total < comp[0].first && low[viz] < prof[centroid] ) total += sub[viz];
		else dfs( viz, centroid, comp[1].first, comp[1].second, 0 );
	}
	if( total < comp[0].first ) return vector<int>(n, 0);

	dfs( 0, 0, comp[0].first, comp[0].second, 1 );

	return resp;
}

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