Submission #830251

#	Time	Username	Problem	Language	Result	Execution time	Memory
830251		de_sousa	Split the Attractions (IOI19_split)	C++17	100 / 100	105 ms	25908 KiB

split

#include "split.h"
#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;

#define all(x) begin(x),end(x)

template<class T>
    using iset=tree<T,null_type,less<T>,rb_tree_tag,tree_order_statistics_node_update>;

struct DSU
{
    int n;
    vector<int> ids;
    vector<int> size;
    DSU(int n):
	n(n),
	ids(vector<int>(n)),
	size(vector<int>(n))
    {
	for(int i=0;i<n;i++){
	    ids[i]=i;
	    size[i]=1;
	}
    }

    int id(int x)
    {
	if(ids[x]==x)return x;
	else return ids[x]=id(ids[x]);
    }
    
    void join(int x,int y)
    {
	x=id(x);
	y=id(y);
	if(x==y)return;
	if(size[x]<size[y])swap(x,y);
	size[x]+=size[y];
	ids[y]=x;
    }
};


vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    int m=p.size();
    map<int,int> f;
    
    {
	vector<array<int,2>> tmp{{a,1},{b,2},{c,3}};
	sort(all(tmp));
		
	a=tmp[0][0];
	b=tmp[1][0];
	c=tmp[2][0];

	f[1]=tmp[0][1];
	f[2]=tmp[1][1];
	f[3]=tmp[2][1];
    }

    // cout<<f[1]<<' '<<f[2]<<' '<<f[3]<<'\n';
    
    vector<vector<int>> g(n);
    for(int i=0;i<m;i++){
	g[p[i]].push_back(q[i]);
	g[q[i]].push_back(p[i]);
    }

    vector<vector<int>> treeg(n);
    vector<int> visited(n);
    auto dfs=[&](auto&&dfs,int x)->void
    {
	visited[x]=true;
	for(int y:g[x]){
	    if(visited[y])continue;
	    treeg[x].push_back(y);
	    treeg[y].push_back(x);
	    dfs(dfs,y);
	}
    };
    dfs(dfs,0);
    
    vector<int> subtreeSize(n);
    auto getSubtreeSize=[&](auto&&getSubtreeSize,int x,int p)->void
    {
	subtreeSize[x]=1;
	for(int y:treeg[x]){
	    if(y==p)continue;
	    getSubtreeSize(getSubtreeSize,y,x);
	    subtreeSize[x]+=subtreeSize[y];
	}
    };

    auto getCentroid=[&](auto&&getCentroid,int x,int p)->int
    {
	for(int y:treeg[x]){
	    if(y==p)continue;
	    if(subtreeSize[y]>n/2)return getCentroid(getCentroid,y,x);
	}
	return x;
    };

    getSubtreeSize(getSubtreeSize,0,-1);
    int centroid=getCentroid(getCentroid,0,-1);
    getSubtreeSize(getSubtreeSize,centroid,-1);

    int curChild;
    DSU dsu(n);
    auto joinSubtree=[&](auto&&joinSubtree,int x,int p)->void
    {
	for(int y:treeg[x]){
	    if(y==p)continue;
	    dsu.join(y,curChild);
	    joinSubtree(joinSubtree,y,x);
	}
    };
    
    for(int y:treeg[centroid]){
	curChild=y;
	joinSubtree(joinSubtree,y,centroid);
    }

    int fixedID=-1;
    for(int y:treeg[centroid]){
	int x=dsu.id(y);
	if(dsu.size[x]>=a){
	    fixedID=x;
	    break;
	}
    }
    
    if(fixedID==-1)for(int i=0;i<m;i++){
	int x=dsu.id(p[i]);
	int y=dsu.id(q[i]);
	if(x==centroid)continue;
	if(y==centroid)continue;
	dsu.join(x,y);
	x=dsu.id(x);
	if(dsu.size[x]>=a){
	    fixedID=x;
	    break;
	}
    }    
    
    
    if(fixedID==-1){
	return vector<int>(n);
    }

    int otherID=-1;
    for(int i=0;i<n;i++){
	int x=dsu.id(i);
	if(x==fixedID)continue;
	otherID=x;
	break;
    }
    assert(otherID!=-1);
    
    for(int i=0;i<n;i++){
	int x=dsu.id(i);
	if(x==fixedID)continue;
	dsu.join(otherID,i);
    }
    dsu.join(otherID,centroid);

    vector<int> sol(n,3);
    auto occupy=[&](auto&&occupy,int x,int p,int&big,int val)
    {
	if(big==0)return;
	sol[x]=val;
	big--;
	for(int y:g[x]){
	    if(sol[y]!=3)continue;
	    if(dsu.id(y)!=dsu.id(x))continue;
	    occupy(occupy,y,x,big,val);
	}
    };

    occupy(occupy,fixedID,-1,a,1);
    occupy(occupy,otherID,-1,b,2);
    
    for(int&i:sol){
	i=f[i];
    }
    return sol;
    
}

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