Submission #423647

#TimeUsernameProblemLanguageResultExecution timeMemory
423647JasiekstrzSplit the Attractions (IOI19_split)C++17
40 / 100
153 ms19688 KiB
#include "split.h"
#include<bits/stdc++.h>
#define fi first
#define se second
using namespace std;
const int N=1e5;
const int INF=1e9+7;
vector<int> ans;
vector<int> e[N+10];
int dp[N+10][2];
int siz[N+10];
void dfs_tree(int x,int p,int a,int b)
{
	siz[x]=1;
	for(auto v:e[x])
	{
		if(v!=p)
		{
			dfs_tree(v,x,a,b);
			siz[x]+=siz[v];
		}
	}
	for(int c:{0,1})
	{
		dp[x][c]=-INF;
		if(siz[x]>=a)
			dp[x][c]=0;
		for(auto v:e[x])
		{
			if(v!=p)
				dp[x][c]=max(dp[x][c],dp[v][c]+siz[x]-siz[v]);
		}
		swap(a,b);
	}
	return;
}
void solve_all(int x,int p,bool c,array<int,2> &s,array<int,2> &id)
{
	if(s[c]==0)
		return;
	ans[x]=id[c];
	s[c]--;
	for(auto v:e[x])
	{
		if(v!=p)
			solve_all(v,x,c,s,id);
	}
	return;
}
void solve_tree(int x,int p,bool c,array<int,2> &s,array<int,2> &id)
{
	if(dp[x][c]==0)
	{
		solve_all(x,p,c,s,id);
		return;
	}
	if(s[!c]>0)
	{
		ans[x]=id[!c];
		s[!c]--;
	}
	bool g=false;
	for(auto v:e[x])
	{
		if(v!=p && !g && dp[x][c]==dp[v][c]+siz[x]-siz[v])
		{
			g=true;
			solve_tree(v,x,c,s,id);
		}
		else if(v!=p)
			solve_all(v,x,!c,s,id);
	}
	return;
}
void dfsok(int x,int &si,int c)
{
	if(si==0)
		return;
	if(ans[x]!=0)
		return;
	ans[x]=c;
	si--;
	for(auto v:e[x])
		dfsok(v,si,c);
	return;
}
vector<int> find_split(int n,int a,int b,int c,vector<int> p,vector<int> q)
{
	int m=p.size();
	ans.resize(n);
	fill(ans.begin(),ans.end(),0);
	for(int i=0;i<m;i++)
	{
		e[p[i]].push_back(q[i]);
		e[q[i]].push_back(p[i]);
	}
	vector<pair<int,int>> s={{a,1},{b,2},{c,3}};
	sort(s.begin(),s.end());
	if(m==n-1)
	{
		dfs_tree(0,-1,s[0].fi,s[1].fi);
		bool c=false;
		if(dp[0][c]<s[!c].fi)
			c=!c;
		if(dp[0][c]<s[!c].fi)
			return ans;
		array<int,2> tmp={s[0].fi,s[1].fi};
		array<int,2> id={s[0].se,s[1].se};
		solve_tree(0,-1,c,tmp,id);
		for(auto &v:ans)
		{
			if(v==0)
				v=s[2].se;
		}
		return ans;
	}
	int si=s[1].fi;
	dfsok(0,si,s[1].se);
	si=s[0].fi;
	for(int i=0;i<n;i++)
		dfsok(i,si,s[0].se);
	for(auto &v:ans)
	{
		if(v==0)
			v=s[2].se;
	}
	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...