Submission #297475

#	Time	Username	Problem	Language	Result	Execution time	Memory
297475		amoo_safar	Split the Attractions (IOI19_split)	C++17	100 / 100	206 ms	27768 KiB

split

#include "split.h"


#include <bits/stdc++.h>

#define pb push_back
#define F first
#define S second
#define all(x) x.begin(), x.end();

using namespace std;

typedef long long ll;

const int N = 2e5 + 10;

int cm[] = {0, 1, 2, 3};

vector<int> G[N], T[N];
int ans[N];
int mk[N];
int sub[N];

void Prep(int u){
	mk[u] = 1;
	sub[u] = 1;
	for(auto adj : G[u]){
		if(!mk[adj]){
			Prep(adj);
			sub[u] += sub[adj];
			T[u].pb(adj);
			T[adj].pb(u);
		}
	}
}

void DFS(int u, int p = -1){
	sub[u] = 1;
	for(auto adj : T[u]){
		if(adj != p){
			DFS(adj, u);
			sub[u] += sub[adj];
		}
	}
}

int Find(int u, int rq, int col){
	rq --;
	ans[u] = col;
	mk[u] = 1;
	if(!rq) return 0;
	for(auto adj : T[u]){
		if(mk[adj])
			continue;
		rq = Find(adj, rq, col);
		if(!rq) break;
	}
	return rq;
}

////////////////
vector<int> vis;
int Find2(int u, int rq, int col){
	vis.pb(u);
	rq --;
	ans[u] = col;
	mk[u] = 1;
	if(!rq) return 0;
	for(auto adj : T[u]){
		if(mk[adj])
			continue;
		rq = Find2(adj, rq, col);
		if(!rq) break;
	}
	if(!rq) return 0;
	for(auto adj : G[u]){
		if(mk[adj])
			continue;
		rq = Find2(adj, rq, col);
		if(!rq) break;
	}
	return rq;
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	if(a > b) swap(a, b), swap(cm[1], cm[2]);
	if(a > c) swap(a, c), swap(cm[1], cm[3]);
	if(b > c) swap(b, c), swap(cm[2], cm[3]);
	int m = p.size();

	for(int i = 0; i < m; i++){
		G[p[i]].pb(q[i]);
		G[q[i]].pb(p[i]);
	}
	Prep(0);
	int cen = 0, fnd = false;
	while(!fnd){
		fnd = true;
		for(auto adj : T[cen]){
			if(sub[adj] < sub[cen] && sub[adj] + sub[adj] >= n){
				fnd = false;
				cen = adj;
				break;
			}
		}
	}
	DFS(cen);
	memset(mk, 0, sizeof mk);

	vector<int> split(n, 0);
	for(auto son : T[cen]){
		if(sub[son] >= a){
			mk[cen] = 1;
			assert( Find(son, a, cm[1]) == 0 );
			mk[cen] = 0;
			assert(	Find(cen, b, cm[2]) == 0 );
			for(int i = 0; i < n; i++){
				if(c && ans[i] == 0){
					ans[i] = cm[3];
					c --;
				}
				split[i] = ans[i];
			}
			return split;
		}
	}
	/////////////////////
	memset(mk, 0, sizeof mk);
	
	mk[cen] = 1;

	for(auto son : T[cen]){
		if(mk[son]) continue;
		vis.clear();
		int rq = Find2(son, a, cm[1]);
		if(rq) continue;
		memset(mk, 0, sizeof mk);
		for(auto x : vis){
			ans[x] = cm[1];
			mk[x] = 1;
		}
		mk[cen] = 0;
		assert(Find(cen, b, cm[2]) == 0);
		for(int i = 0; i < n; i++){
			if(c && ans[i] == 0){
				ans[i] = cm[3];
				c --;
			}
			split[i] = ans[i];
		}
		return split;
	}
	return split;
}

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