Submission #244586

#TimeUsernameProblemLanguageResultExecution timeMemory
244586tictaccatSplit the Attractions (IOI19_split)C++14
100 / 100
202 ms40964 KiB
#include "split.h"
#include <bits/stdc++.h>

using namespace std;

const int mx = 300000;
int sz[mx], sz2[mx];
bool explored[mx], explored2[mx], chosen[mx], A[mx], explored3[mx];
vector<int> mark(mx, 0), res, tbl(3);
vector<vector<int>> adj(mx), ch(mx), adj2(mx);
int n, cent, a, b, c, repA, repB;
bool valid;

void dfs(int u) {
	explored[u] = true;
	sz[u] = 1;
	bool works = true;
	for (int v: adj[u]) {
		if (explored[v]) continue;
		//cout << u << " " << v << "\n";
		ch[u].push_back(v);
		dfs(v);
		if (sz[v] > n/2) works = false;
		sz[u] += sz[v];
	}
	if (n - sz[u] > n/2) works = false;
	if (works) cent = u;
}

void dfs2(int u) {
	for (int v: ch[u]) {
		//cout << v << "\n";
		mark[v] = mark[u];
		dfs2(v);
	}
}

void dfs3(int u, vector<int>& comp) {
	explored2[u] = true;
	for (int v: adj2[u]) {
		if (explored2[v]) continue;
		comp.push_back(v);
		dfs3(v, comp);
	}
}

void dfs4(int u, int& cnt, bool inA) {
	explored3[u] = true;
	cnt++;
	if (inA && cnt <= a) {
		res[u] = tbl[0];
	}
	else if (!inA && cnt <= b) {
	//	cout << u << "\n";
		res[u] = tbl[1];
	}
	for (int v: adj[u]) {
		if (explored3[v] || A[v] != inA) continue;
		dfs4(v, cnt, inA);
	}
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	::n = n;
	vector<int> arr = {a,b,c};
	tbl = {1,2,3};
	sort(tbl.begin(),tbl.end(),[&arr](int i, int j){return arr[i-1] < arr[j-1];});
	sort(arr.begin(),arr.end());
	a = arr[0];
	b = arr[1];
	c = arr[2];
	::a = a;
	::b = b;
	::c = c;
	//cout << tbl[0] << " " << tbl[1] << " " << tbl[2] << "\n";
	//assume a <= b <= c
	for (int i = 0; i < (int)p.size(); i++) {
		adj[p[i]].push_back(q[i]);
		adj[q[i]].push_back(p[i]);
	}
	dfs(2);

	//cout << cent << "\n";

	mark[cent] = -1;
	sz2[0] = n - sz[cent];
	for (int i = 0; i < (int)ch[cent].size(); i++) {
		int u = ch[cent][i];
		//cout << u << "\n";
		mark[u] = i+1;
		sz2[mark[u]] = sz[u];
		dfs2(u);
	}

	//cout << mark[7] << "\n";

	for (int i = 0; i < (int)p.size(); i++) {
		if (p[i] == cent || q[i] == cent) continue;
	//	cout << p[i] << " " << q[i] << "\n";
		if (mark[p[i]] != mark[q[i]]) {
			adj2[mark[p[i]]].push_back(mark[q[i]]);
			adj2[mark[q[i]]].push_back(mark[p[i]]);
			//cout << mark[p[i]] << " " << mark[q[i]] << "\n";
		}
	}

	for (int i = 0; i <= (int)ch[cent].size(); i++) {
		vector<int> comp;
		comp.push_back(i);
		dfs3(i, comp);
		int x = 0;
		for (int c: comp) {
			x += sz2[c];
		}
		//cout << x << "\n";
		if (x >= a) {
			sort(comp.begin(),comp.end(),[](int i, int j){return sz2[i] > sz2[j];});
			while (x - sz2[comp.back()] >= a) {
				x -= sz2[comp.back()];
				comp.pop_back();
			}
			//cout << comp[0] << "\n";
			for (int c: comp) chosen[c] = true;
			for (int i = 0; i < n; i++) {
				if (i == cent) continue;
				A[i] = chosen[mark[i]];
				if (A[i]) {
					//cout << i << " ";
					repA = i;
				}
				else repB = i;
			}
			//cout << "hooray\n";
			valid = true;
			break;
		}
	}

	if (valid) {
		res.assign(n,tbl[2]);
		int cnt = 0;
		//cout << repA << " " << repB << "\n";
		dfs4(repA, cnt, true);
		cnt = 0;
		dfs4(repB, cnt, false);
	}

	if (!valid) {
		res.assign(n,0);
	}

	return res;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...