Submission #1047540

#TimeUsernameProblemLanguageResultExecution timeMemory
1047540vjudge1Split the Attractions (IOI19_split)C++17
40 / 100
436 ms1048576 KiB
#include "split.h"

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

const int lim=200100;

#define pb push_back

vector<int>res;
int sszz;
int too;
bool vis[lim];

vector<int>v[lim];

void dfs(int node){
	if(!sszz||vis[node])return;
	vis[node]=1;
	res[node]=too;
	sszz--;
	for(int j:v[node]){
		dfs(j);
	}
}

int AA,BB,CC;

void dfs1(int node){
	if(vis[node])return;
	vis[node]=1;
	if(AA){
		res[node]=1;
		AA--;
	}else if(BB){
		res[node]=2;
		BB--;
	}else{
		res[node]=3;
		CC--;
	}
	for(int j:v[node]){
		dfs1(j);
	}
}

int sz[lim],pr[lim];

int sbs(int node,int par){
	pr[node]=par;
	sz[node]=1;
	for(int j:v[node]){
		if(j==par)continue;
		sz[node]+=sbs(j,node);
	}
	return sz[node];
}

int need[4];
int least,mid,maxi;

void ddfs(int node,int par,int ty){
	if(need[ty]){
		res[node]=ty;
		need[ty]--;
	}
	for(int j:v[node]){
		if(j==par)continue;
		ddfs(j,node,ty);
	}
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	res=vector<int>(n);
	int m=p.size();
	for(int i=0;i<m;i++){
		v[p[i]].pb(q[i]);
		v[q[i]].pb(p[i]);
	}
	if(a==1){
		if(b<c){
			fill(res.begin(),res.end(),3);
			too=2;
		}else{
			fill(res.begin(),res.end(),2);
			too=3;
		}
		sszz=min(b,c);
		dfs(0);
		for(int i=0;i<n;i++){
			if(res[i]!=too){
				res[i]=1;
				break;
			}
		}
		return res;
	}
	bool ok=1;
	for(int i=0;i<n;i++){
		if(2<v[i].size()){
			ok=0;
			break;
		}
	}
	AA=a,BB=b,CC=c;
	if(ok){
		int root=0;
		for(int i=0;i<n;i++){
			if(v[i].size()==1){
				root=i;
			}
		}
		dfs1(root);
	}else{
		need[1]=a,need[2]=b,need[3]=c;
		sbs(0,0);
		if(a<b&&a<c){
			least=1;
			if(b<c){
				mid=2;
				maxi=3;
			}else{
				mid=3;
				maxi=2;
			}
		}
		else if(b<c){
			least=2;
			if(a<c){
				mid=1;
				maxi=3;
			}else{
				mid=3;
				maxi=1;
			}
		}
		else{
			least=3;
			if(a<b){
				mid=1;
				maxi=2;
			}else{
				mid=2;
				maxi=1;
			}
		}
		for(int i=0;i<n;i++){
			res[i]=maxi;
		}
		int mini=-1,ss=INT_MAX;
		for(int i=0;i<n;i++){
			if(need[least]<=sz[i]&&sz[i]<ss){
				ss=sz[i];
				mini=i;
			}
		}
		if(need[mid]<=n-ss){
			ddfs(mini,pr[mini],least);
			ddfs(pr[mini],mini,mid);
			return res;
		}
		mini=-1,ss=INT_MAX;
		for(int i=0;i<n;i++){
			if(need[mid]<=sz[i]&&sz[i]<ss){
				ss=sz[i];
				mini=i;
			}
		}
		if(need[least]<=n-ss){
			ddfs(mini,pr[mini],mid);
			ddfs(pr[mini],mini,least);
			return res;
		}
		for(int i=0;i<n;i++){
			res[i]=0;
		}
		return res;
	}
	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...