Submission #285612

#	Time	Username	Problem	Language	Result	Execution time	Memory
285612		MKopchev	Split the Attractions (IOI19_split)	C++14	100 / 100	254 ms	28652 KiB

split

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

using namespace std;

const int nmax=2e5+42;

pair<int,int> edges[nmax];

int n,m;

pair<int/*size*/,int/*colour*/> given[3];

vector<int> outp;

vector< pair<int,int> > extra,create_first;

int SZ[nmax],parent[nmax];
int root(int node)
{
    if(parent[node]==node)return node;
    parent[node]=root(parent[node]);
    return parent[node];
}
vector<int> adj[nmax],all_adj[nmax];

int centroid;

void dfs(int node,int par)
{
    SZ[node]=1;

    for(auto k:adj[node])
        if(k!=par)
        {
            dfs(k,node);
            SZ[node]+=SZ[k];
        }
}
int find_centroid()
{
    dfs(0,0);

    int ret=0;

    while(SZ[ret]>SZ[0]/2)
    {
        bool stop=1;

        for(auto k:adj[ret])
            if(SZ[ret]>SZ[k]&&SZ[k]>SZ[0]/2)
            {
                stop=0;
                ret=k;
                break;
            }

        if(stop)break;
    }
    return ret;
}

void dfs_create_adj(int node)
{
    //cout<<"dfs_create_adj "<<node<<endl;

    if(given[0].first==0)return;

    if(outp[node])return;

    given[0].first--;
    outp[node]=given[0].second;

    //cout<<node<<" -> "<<outp[node]<<endl;

    for(auto k:adj[node])
        dfs_create_adj(k);
}

void dfs_create_all(int node)
{
    //cout<<"dfs_create_all "<<node<<endl;

    if(given[1].first==0)return;

    if(outp[node])return;

    given[1].first--;
    outp[node]=given[1].second;

    //cout<<node<<" -> "<<outp[node]<<endl;

    for(auto k:all_adj[node])
        dfs_create_all(k);
}

void fill_dumb()
{
    for(auto &k:outp)
        if(k==0)
        {
            k=given[2].second;
            given[2].first--;
        }
    assert(given[0].first==0&&given[1].first==0&&given[2].first==0);
}

void push(vector< pair<int,int> > cur)
{
    if(given[0].first==1)
    {
        dfs_create_adj((centroid==0?1:0));
        dfs_create_all(centroid);
        fill_dumb();

        return;
    }

    for(auto k:cur)
        if(k.first!=centroid&&k.second!=centroid&&root(k.first)!=root(k.second))
        {
            adj[k.first].push_back(k.second);
            adj[k.second].push_back(k.first);

            //cout<<"pushed "<<k.first<<" "<<k.second<<endl;

            SZ[root(k.first)]+=SZ[root(k.second)];
            parent[root(k.second)]=root(k.first);

            if(SZ[root(k.first)]>=given[0].first)
            {
                dfs_create_adj(k.first);
                dfs_create_all(centroid);
                fill_dumb();

                return;
            }
        }
}
vector<int> find_split(int N, int a, int b, int c, vector<int> p, vector<int> q)
{
    n=N;

    outp={};

    for(int i=0;i<n;i++)
    {
        outp.push_back(0);
        all_adj[i]={};
    }

	m=p.size();

	given[0]={a,1};
	given[1]={b,2};
	given[2]={c,3};

	sort(given,given+3);

	for(int i=0;i<m;i++)
    {
        edges[i]={p[i],q[i]};

        all_adj[p[i]].push_back(q[i]);
        all_adj[q[i]].push_back(p[i]);
    }

    for(int i=0;i<n;i++)
    {
        SZ[i]=1;
        parent[i]=i;
    }

    for(int i=0;i<m;i++)
        if(root(p[i])!=root(q[i]))
        {
            SZ[root(p[i])]+=SZ[root(q[i])];
            parent[root(q[i])]=root(p[i]);

            adj[p[i]].push_back(q[i]);
            adj[q[i]].push_back(p[i]);

            create_first.push_back({p[i],q[i]});
        }
        else extra.push_back({p[i],q[i]});

    centroid=find_centroid();

    //cout<<"centroid= "<<centroid<<endl;

    for(int i=0;i<n;i++)
    {
        adj[i]={};
        SZ[i]=1;
        parent[i]=i;
    }

    push(create_first);

    if(outp[0])return outp;

    push(extra);

    return outp;
}

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