Submission #1031313

#	Time	Username	Problem	Language	Result	Execution time	Memory
1031313		Lalic	Split the Attractions (IOI19_split)	C++17	40 / 100	99 ms	23464 KiB

split

#include "split.h"
#include <bits/stdc++.h>
using namespace std;

#define fi first
#define se second
#define pb push_back
#define all(x) x.begin(), x.end()
#define allr(x) x.rbegin(), x.rend()
#define mp make_pair

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef complex<double> cd;

const int MAXN = 1e5+10;

vector<int> proc[MAXN], adj[MAXN];
int cor[MAXN], mark[MAXN], sz[MAXN], N;

void build(int ver){
	mark[ver]=1;
	
	for(auto u : proc[ver]){
		if(mark[u]) continue;
		adj[ver].pb(u);
		adj[u].pb(ver);
		build(u);
	}
}

int getSubtreeSizes(int ver, int prev=-1){
	sz[ver]=1;
	
	for(auto u : adj[ver]){
		if(u==prev) continue;
		sz[ver]+=getSubtreeSizes(u, ver);
	}
	
	return sz[ver];
}

void setColor(int ver, int prev, int& qnt, int c){
	if(!qnt) return;
	cor[ver]=c;
	qnt--;
	if(!qnt) return;
	
	for(auto u : adj[ver]){
		if(u==prev) continue;
		setColor(u, ver, qnt, c);
	}
}

bool calc(int ver, int prev, int big, int small, int bc, int sc){
	pii mx={N-sz[ver], prev};
	for(auto u : adj[ver]){
		if(u==prev) continue;
		mx=max(mx, {sz[u], u});
	}
	
	if(mx.fi>=big && N-mx.fi>=small){
		int auxsmall=small-1, auxbig=big;
		cor[ver]=sc;
		
		//~ cout << "bigcolor: " << bc << "\nsmallcolor: " << sc << "\n";
		
		for(auto u : adj[ver]){
			if(u==mx.se) setColor(u, ver, auxbig, bc);
			else setColor(u, ver, auxsmall, sc);
		}
		
		//~ cout << auxsmall << "\n";
		
		return 1;
	}
	
	for(auto u : adj[ver]){
		if(u==prev) continue;
		if(calc(u, ver, big, small, bc, sc)) return 1;
	}
	
	return 0;	
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	N=n;
	
	int m=(int)p.size();
	
	for(int i=0;i<m;i++){
		proc[p[i]].pb(q[i]);
		proc[q[i]].pb(p[i]);
	}
	
	build(0);
	getSubtreeSizes(0);
	
	int big, small, bigcolor, smallcolor;
	vector<pii> temp={{a, 1}, {b, 2}, {c, 3}};
	sort(all(temp));
	
	//~ cout << "temp:\n";
	//~ for(auto u : temp) cout << u.fi << " " << u.se << "\n";
	
	big=temp[1].fi; bigcolor=temp[1].se;
	small=temp[0].fi; smallcolor=temp[0].se;
	
	if(!calc(0, -1, big, small, bigcolor, smallcolor)){
		vector<int> resp;
		for(int i=0;i<n;i++) resp.pb(0);
		return resp;
	}
	
	for(int i=0;i<n;i++)
		if(!cor[i]) cor[i]=temp[2].se;
	
	vector<int> ans(n);
	for(int i=0;i<n;i++) ans[i]=cor[i];
	
	return ans;
}

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