Submission #506197

#	Time	Username	Problem	Language	Result	Execution time	Memory
506197		MilosMilutinovic	Split the Attractions (IOI19_split)	C++14	100 / 100	131 ms	32708 KiB

split

#include "split.h"
#include <bits/stdc++.h>

using namespace std;

using vi=vector<int>;

const int nax=2e5+5;

int n;
int a, b, c;
int ord[4];

vi graf[nax];
vi tree[nax];

int sz[nax];

int komp[nax];

void dfs_sz(int u, int p)
{
	sz[u]=1;
	for (int v:tree[u])
	{
		if (v==p) continue;
		dfs_sz(v, u);
		sz[u]+=sz[v];
	}
}

int dfsf(int u, int p)
{
	for (int v:tree[u])
	{
		if (v==p) continue;
		if (sz[v]*2>=n)
			return dfsf(v, u);
	}
	return u;
}

int centroid()
{
	dfs_sz(0, 0);
	return dfsf(0, 0);
}

void dfs_down(int u, int p, vi& ans)
{
	if (a==0) return;
	a--;
	//assert(ans[u]==0);
	ans[u]=ord[1];
	for (int v:tree[u])
	{
		if (v==p) continue;
		dfs_down(v, u, ans);
	}
}

void dfs_fill2(int u, int p, vi& ans)
{
	if (b==0) return;
	b--;
	assert(ans[u]==0);
	ans[u]=ord[2];
	for (int v:tree[u])
	{
		if (v==p||ans[v]!=0) continue;
		dfs_fill2(v, u, ans);
	}
}

bool was[nax];

void dfs_tree(int u, int p)
{
	was[u]=true;
	for (int v:graf[u])
	{
		if (was[v]) continue;
		tree[u].push_back(v);
		tree[v].push_back(u);
		dfs_tree(v, u);
	}
}

void dfs_par(int u, int p, int c)
{
	komp[u] = c;
	for (int v:tree[u])
	{
		if (v==p) continue;
		dfs_par(v, u, c);
	}
}

vi ct[nax];
vi ls;

int vis[nax];

int tot;

void dfs_fin(int u)
{
	ls.push_back(u);
	vis[u]=1;
	tot+=sz[u];
	if (tot>=a) return;
	for (int v:ct[u])
	{
		if (!vis[v]) dfs_fin(v);
		if (tot>=a) return;
	}
}

vi find_split(int N, int A, int B, int C, vi u, vi v)
{
	n=N; a=A; b=B; c=C;
	ord[1]=1;
	ord[2]=2;
	ord[3]=3;
	if (a>c)
	{
		swap(a,c);
		swap(ord[1],ord[3]);
	}
	if (b>c)
	{
		swap(b,c);
		swap(ord[2],ord[3]);
	}
	if (a>b)
	{
		swap(a,b);
		swap(ord[1],ord[2]);
	}
	int m=v.size();
	for (int i=0; i<m; i++)
	{
		graf[u[i]].push_back(v[i]);
		graf[v[i]].push_back(u[i]);
	}
	if (m==n-1)
	{
		for (int i=0;i<m;i++)
		{
			tree[u[i]].push_back(v[i]);
			tree[v[i]].push_back(u[i]);
		}
		C=c;
		int c=centroid();
		dfs_sz(c, c);
		for(int v:graf[c])
		{
			if (sz[v]>=a)
			{
				vi ans(n);
				dfs_down(v, c, ans);
				//assert(sz[c]-sz[v]>=b);
				dfs_fill2(c, v, ans);
				for (int i=0; i<n; i++)
				{
					if (ans[i]==0)
					{
						C--;
						ans[i]=ord[3];
					}
				}
				return ans;
			}
		}
		vi ans(n);
		return ans;
	}
	dfs_tree(0, 0);
	int cen=centroid();
	dfs_sz(cen, cen);
	for (int v:tree[cen])
	{
		if (sz[v]>=a)
		{
			vi ans(n);
			dfs_down(v, cen, ans);
			dfs_fill2(cen, v, ans);
			for (int i=0; i<n; i++)
			{
				if (ans[i]==0)
				{
					ans[i]=ord[3];
				}
			}
			return ans;
		}
	}
	for (int v:tree[cen])
	{
		dfs_par(v,cen,v);
	}
	for (int i=0; i<m; i++)
	{
		if (u[i]!=cen && v[i]!=cen && komp[u[i]] != komp[v[i]])
		{
			ct[komp[u[i]]].push_back(komp[v[i]]);
			ct[komp[v[i]]].push_back(komp[u[i]]);
		}
	}
	for (int v:tree[cen])
	{
		if (!vis[v])
		{
			ls.clear();
			tot=0;
			dfs_fin(v);
			if (tot>=a)
			{
				vi ans(n);
				for (int who:ls)
				{
					dfs_down(who, cen, ans);
				}
				bool ok=false;
				for (int i=0; i<n; i++)
					if (ans[i]==ord[1])
						ok=true;
				assert(ok);
				dfs_fill2(cen, cen, ans);
				for (int i=0; i<n; i++)
				{
					if (ans[i]==0)
					{
						ans[i]=ord[3];
					}
				}
				return ans;
			}
		}
	}
	vi ans(n);
	return ans;
}

/*
9 10
4 2 3
0 1
0 2
0 3
0 4
0 6
0 8
1 7
3 7
4 5
5 6
*/

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