Submission #589575

#	Time	Username	Problem	Language	Result	Execution time	Memory
589575		farhan132	Split the Attractions (IOI19_split)	C++17	100 / 100	188 ms	29444 KiB

split

#include "split.h"
#include <bits/stdc++.h>
 
using namespace std;
 
typedef int ll;
typedef pair<ll , ll> ii;
 
#define ff first
#define ss second
#define pb push_back
#define in insert

const ll N = 2e5 + 5;

ll n;

struct DSU{
	ll par[N], sz[N];
	void start(){
		for(ll i = 0; i < N; i++){
			par[i] = i; sz[i] = 1;
		}
		return;
	}
	ll find(ll a){
		if(a == par[a]) return par[a];
		return par[a] = find(par[a]);
	}
	void merge(ll x, ll y){
		x = find(x); y = find(y);
		if(x == y) return;
		if(sz[x] > sz[y]) swap(x, y);
		sz[y] += sz[x];
		par[x] = y;
		return;
	}
};

vector < ll > v[N], g[N];
ll sub[N];
ll up,down;

void dfs(ll s, ll p){
	sub[s] = 1;
	for(auto u : v[s]){
		if(u - p){
			dfs(u, s);
			sub[s] += sub[u];
		}
	}
	return;
}
ll centroid(ll s, ll p){
	for(auto u : v[s]){
		if(u - p){
			if(sub[u] > n/2){
				return centroid(u, s);
			}
		}
	}
	return s;
}
vector < ll > res;
ll cen;

void color(ll s, ll &rem, ll col){
	if(s == cen || res[s] == col) return;
	if(rem == 0) return;
	res[s] = col; rem--;
	for(auto u : v[s]){
			color(u, rem, col);
	}
	return;
}
ll wut[N];

void full(ll s, ll p, ll x){
	wut[s] = x;
	for(auto u : v[s]){
		if(u - p){
			full(u, s, x);
		}
	}
}

vector<int> find_split(int _n, int a, int b, int c, vector<int> p, vector<int> q) {
	vector < ii > all;
	all.pb({a, 1});
	all.pb({b, 2});
	all.pb({c, 3});
	sort(all.begin(), all.end());
	down = all[0].ff; up = all[1].ff; n = _n;

	ll m = p.size();

	vector < ii > extra;

	DSU T; T.start();

	for(ll i = 0; i < m; i++){
		ll x = p[i], y = q[i];
		g[x].pb(y);
		g[y].pb(x);
		if(T.find(x) == T.find(y)){
			extra.pb({x, y});
			continue;
		}
		T.merge(x, y);
		v[x].pb(y);
		v[y].pb(x);
	}
	res.resize(n, 0);

	dfs(0, -1);
	cen = centroid(0, -1);
	dfs(cen, -1);
	ll f = -1;
	for(auto u : v[cen]){
		if(sub[u] >= down){
			f = u;
			break;
		}
	}
	if(f != -1){
		res.resize(0);
		res.resize(n, all[2].ss);

		color(f, down, all[0].ss);

		up--;
		res[cen] = all[1].ss;
		for(auto u : v[cen]){
			if(u == f) continue;
			color(u, up, all[1].ss);
		}
		return res;
	}

	DSU final;

	final.start();

	for(auto u : v[cen]){
		final.sz[u] = sub[u];
		full(u, cen, u);
	}
	vector < ii > last_hope;
	for(auto [x, y] : extra){
		if(x == cen || y == cen) continue;
		if(final.find(wut[x]) == final.find(wut[y])) continue;
		ll X = final.find(wut[x]); ll Y = final.find(wut[y]);
		if(final.sz[X] + final.sz[Y] > n/2){
			 last_hope.pb({x, y}); continue;
		}
		ll kk = final.sz[X] + final.sz[Y];
		final.merge(X, Y);
		v[x].pb(y);
		v[y].pb(x);
		//assert(final.sz[final.find(wut[x])] == kk && kk == final.sz[final.find(wut[y])]);
 	}
 	f = -1;
 	for(auto u : v[cen]){
		if(final.sz[final.find(u)] >= down){
			f = u;
			break;
		}
	}
	if(f != -1){
		res.resize(0);
		res.resize(n, all[2].ss);

		color(f, down, all[0].ss);
		up--;
		res[cen] = all[1].ss;
		for(auto u : v[cen]){
			if(final.find(u) == final.find(f)) continue;
			color(u, up, all[1].ss);
		}
		return res;
	}

	for(auto [x, y] : last_hope){
		ll X = final.find(wut[x]); ll Y = final.find(wut[y]);
		if(n - final.sz[X] + final.sz[Y] >= all[1].ff){
			final.merge(X, Y);
			v[x].pb(y);
			v[y].pb(x);
			break;
		}
	}
	f = -1;
 	for(auto u : v[cen]){
		if(final.sz[final.find(u)] >= down){
			f = u;
			break;
		}
	}
	if(f != -1){
		res.resize(0);
		res.resize(n, all[2].ss);

		color(f, down, all[0].ss);

		up--;
		res[cen] = all[1].ss;
		for(auto u : v[cen]){
			if(final.find(u) == final.find(f)) continue;
			color(u, up, all[1].ss);
		}
		return res;
	}

	return res;
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:156:6: warning: unused variable 'kk' [-Wunused-variable]
  156 |   ll kk = final.sz[X] + final.sz[Y];
      |      ^~

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