Submission #1043782

#	Time	Username	Problem	Language	Result	Execution time	Memory
1043782		mariaclara	Split the Attractions (IOI19_split)	C++17	100 / 100	87 ms	26576 KiB

split

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

using namespace std;

typedef long long ll;
typedef pair<int,int> pii;
const int MAXN = 1e5+5;
#define all(x) x.begin(), x.end()
#define sz(x) (int)x.size()
#define mk make_pair
#define pb push_back
#define fr first
#define sc second

int n, niv[MAXN], sub[MAXN], back[MAXN], who[MAXN], to_paint, cent;
bool vis[MAXN];
vector<int> edges[MAXN], tree[MAXN], color;
vector<pii> comp;

void calc_sub(int x) {
	sub[x] = vis[x] = 1;
	back[x] = who[x] = x;

	//cout << "at " << x << "\n\n";
	for(int viz : edges[x]) {
		if(!vis[viz]) {
			niv[viz] = niv[x] + 1;
			//cout << "pai " << x << "\n";
			calc_sub(viz);
			sub[x] += sub[viz];
			tree[x].pb(viz);
			tree[viz].pb(x);
			if(niv[back[viz]] < niv[back[x]]) back[x] = back[viz], who[x] = viz;
		}
		else if(niv[viz] < niv[back[x]]) back[x] = viz, who[x] = x;
	}
}

int find_cent(int x) {
	comp.clear();
	vis[x] = 1;

	if(x != 0) comp.pb({n-sub[x], 0});

	for(int viz : edges[x]) {
		if(vis[viz]) continue;
		if(sub[viz] > n/2) return find_cent(viz);
		comp.pb({sub[viz], who[viz]});
	}

	return x;
}

void paint(int x, int c) {
	if(to_paint == 0) return;
	color[x] = c;
	to_paint--;

	for(int viz : tree[x]) if(!color[viz]) paint(viz, c);
}


vector<int> ans(vector<int> v, int a, int b, int c, bool flag) {
	pii aux[3] = {{a,1},{b,2},{c,3}};
	sort(aux, aux+3);
	vector<int> ret;

	for(auto i : v) {
		if(i == 0) ret.pb(aux[2].sc * flag);
		else ret.pb(aux[i-1].sc * flag);
	}

	return ret;
}

vector<int> find_split(int N, int a, int b, int c, vector<int> p, vector<int> q) {
	n = N;
	color.resize(n);

	for(int i = 0; i < sz(p); i++)
		edges[p[i]].pb(q[i]), edges[q[i]].pb(p[i]);

	fill(vis, vis+n, 0);
	calc_sub(0);

	fill(vis, vis+n, 0);
	cent = find_cent(0);
	color[cent] = 2;
	
	//cout << cent << "\n";

	to_paint = min({a,b,c});

	for(auto [s, x] : comp) {
		//cout << s << " " << x << "\n";
		if(s >= to_paint) {
			paint(x, 1);
			to_paint = min({max(a,b), max(b,c), max(c,a)});
			paint(cent, 2);
			return ans(color, a, b, c, 1);
		}
	}

	if(cent == 0) return ans(color, a, b, c, 0);
	int S = comp[0].fr;
	paint(0, 1);
	
	for(auto [s, x] : comp)
		if(x != 0 and niv[back[x]] < niv[cent]) paint(x, 1);

	if(to_paint > 0) return ans(color, a, b, c, 0);

	to_paint = min({max(a,b), max(b,c), max(c,a)});
	paint(cent, 2);
	return ans(color, a, b, c, 1);
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:106:6: warning: unused variable 'S' [-Wunused-variable]
  106 |  int S = comp[0].fr;
      |      ^

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