Submission #538442

#	Time	Username	Problem	Language	Result	Execution time	Memory
538442		jamezzz	Split the Attractions (IOI19_split)	C++17	100 / 100	172 ms	28448 KiB

split

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

#define sf scanf
#define pf printf
#define fi first
#define se second
#define pb push_back
#define sz(x) (int)x.size()
#define all(x) x.begin(),x.end()
typedef pair<int,int> ii;

#define maxn 100005

int a,b,c,mark[3],vis[maxn],pa[maxn],sub[maxn],rem;
int par[maxn],sz[maxn];
vector<int> res,AL[maxn],tree[maxn];
vector<ii> tedge,nedge;//tree edge, non tree edge

int fp(int i){
	return (par[i]==i)?i:par[i]=fp(par[i]);
}

void join(int x,int y){
	x=fp(x),y=fp(y);
	if(x==y)return;
	if(sz[x]>sz[y])swap(x,y);
	par[x]=y;sz[y]+=sz[x];
}

void dfs(int u){
	vis[u]=true;
	sub[u]=1;
	for(int v:AL[u]){
		if(v==pa[u])continue;
		if(vis[v]){
			nedge.pb({u,v});
			continue;
		}
		tedge.pb({u,v});
		tree[u].pb(v);
		tree[v].pb(u);
		pa[v]=u;
		dfs(v);
		sub[u]+=sub[v];
	}
}

void bfs(int rt,int num){
	queue<int> q;
	q.push(rt);
	while(!q.empty()){
		int u=q.front();q.pop();
		for(int v:tree[u]){
			if(res[v]!=0)continue;
			if(num==0)return;
			res[v]=res[rt];
			--num;
			q.push(v);
		}
	}
}

vector<int> find_split(int n,int _a,int _b,int _c,vector<int> p,vector<int> q){
	vector<ii> cols;
	cols.pb({_a,1});
	cols.pb({_b,2});
	cols.pb({_c,3});
	sort(all(cols));
	tie(a,mark[0])=cols[0];
	tie(b,mark[1])=cols[1];
	tie(c,mark[2])=cols[2];
	
	res.resize(n);
	
	for(int i=0;i<sz(p);++i){
		AL[p[i]].pb(q[i]);
		AL[q[i]].pb(p[i]);
	}
	
	memset(pa,-1,sizeof pa);
	dfs(0);
	
	for(auto [x,y]:tedge){
		if(sub[x]>sub[y])swap(x,y);//x is child of y
		int sa=sub[x],sb=n-sub[x];
		if(sa>sb)swap(sa,sb),swap(x,y);
		
		if(sa>=a&&sb>=b){
			res[x]=cols[0].se;
			res[y]=cols[1].se;
			bfs(x,a-1);bfs(y,b-1);
			for(int i=0;i<n;++i){
				if(res[i]==0)res[i]=cols[2].se;
			}
			return res;
		}
	}
	
	int r=-1;
	for(int i=0;i<n;++i){
		bool can=true;
		if(sub[i]<a)can=false;
		for(int j:AL[i]){
			if(sub[j]>sub[i])continue;//parent
			if(sub[j]>=a)can=false;
		}
		if(can){
			r=i;
			break;
		}
	}
	
	for(int i=0;i<n;++i)par[i]=i,sz[i]=1;
	
	for(auto [u,v]:tedge){
		if(u==r||v==r)continue;
		join(u,v);
	}
	
	for(auto [u,v]:nedge){
		if(u==r||v==r)continue;
		tree[u].pb(v);
		tree[v].pb(u);
		join(u,v);
		if(sz[fp(u)]>=a){
			int x=u,y=r;
			res[x]=cols[0].se;
			res[y]=cols[1].se;
			bfs(x,a-1);bfs(y,b-1);
			for(int i=0;i<n;++i){
				if(res[i]==0)res[i]=cols[2].se;
			}
			return res;
		}
	}
	
	return res;
}

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