Submission #256111

#TimeUsernameProblemLanguageResultExecution timeMemory
256111b00n0rpSplit the Attractions (IOI19_split)C++17
18 / 100
164 ms28920 KiB
#include "split.h"
#include <bits/stdc++.h>
using namespace std;

#define ll long long
#define vi vector<int>
#define pb push_back
#define REP(i,n) for(int i = 0; i < n; i++)
#define FOR(i,a,b) for(int i = a; i < b; i++)
#define pii pair<int,int>
#define F first
#define S second
using namespace std;

const int MX = 200005;

vi g[MX],adj[MX];
int par[MX],dep[MX],sz[MX];

bitset<MX> vis;

vi ans;

void dfs(int u,int p,int d = 0){
	// cout << u << " " << p << endl;
	par[u] = p;
	vis[u] = 1;
	dep[u] = d;
	sz[u] = 1;
	for(auto v:g[u]){
		if(vis[v]) continue;
		adj[u].pb(v);
		par[v] = u;
		dfs(v,u,d-1);
		sz[u] += sz[v];
	}
}

void assign(int u,int val){
	ans[u] = val;
	vis[u] = 1;
	for(auto v:g[u]){
		if(!vis[v]) assign(v,val);
	}
}



vi find_split(int n, int a, int b, int c, vi p, vi q) {
	ans.resize(n);
	REP(i,(int)p.size()){
		g[p[i]].pb(q[i]);
		g[q[i]].pb(p[i]);
	}
	int mxd = 0;
	int root = 0;
	REP(i,n){
		mxd = max(mxd,(int)g[i].size());
		if(g[i].size() == 1) root = i;
	}
	dfs(root,root);
	
	if(mxd == 2){
		int cur = root,cnt = 0;
		while(1){
			if(cnt < a) ans[cur] = 1;
			else if(cnt < a+b) ans[cur] = 2;
			else ans[cur] = 3;
			cnt++;
			if(adj[cur].size()) cur = adj[cur][0];
			else break;
		}
		return ans;
	}
	if(a == 1){
		vector<pii> v;
		REP(i,n) v.pb({dep[i],i});
		sort(v.begin(),v.end());
		REP(i,n){
			if(i == 0) ans[v[i].S] = 1;
			else if(i <= b) ans[v[i].S] = 2;
			else ans[v[i].S] = 3;
		}
		return ans;
	}
	vis.reset();
	bool flag = 0;
	REP(i,n){
		root = i;
		if(sz[i] == a){
			vis[par[i]] = 1;
			assign(i,1);
			flag = 1;
			a = 0;
			break;
		}
		if(sz[i] == n-a){
			vis[i] = 1;
			assign(par[i],1);
			flag = 1;
			a = 0;
			break;
		}



		if(sz[i] == b){
			vis[par[i]] = 1;
			assign(i,2);
			flag = 1;
			b = 0;
			break;
		}
		if(sz[i] == n-b){
			vis[i] = 1;
			assign(par[i],2);
			flag = 1;
			b = 0;
			break;
		}



		if(sz[i] == c){
			vis[par[i]] = 1;
			assign(i,3);
			flag = 1;
			c = 0;
			break;
		}
		if(sz[i] == n-c){
			vis[i] = 1;
			assign(par[i],3);
			flag = 1;
			c = 0;
			break;
		}
	}	
	if(!flag) return ans;
	dfs(root,root);
	vector<pii> v;
	REP(i,n){
		if(!ans[i]) v.pb({dep[i],i});
	}
	sort(v.begin(),v.end());
	REP(i,(int)v.size()){
		// cout << v[i].F << " " << v[i].S << endl;
		if(i < a) ans[v[i].S] = 1;
		else if(i < a+b) ans[v[i].S] = 2;
		else ans[v[i].S] = 3;
	}
	return ans;
}
#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...