Submission #1292434

#TimeUsernameProblemLanguageResultExecution timeMemory
1292434MMihalevSplit the Attractions (IOI19_split)C++20
0 / 100
605 ms1114112 KiB
#include<iostream>
#include<vector>
#include<algorithm>
#include<queue>
#include "split.h"
using namespace std;
const int MAX_N=2e5+5;
vector<int>g[MAX_N];
int col[MAX_N];
int deg[MAX_N];
int A,B;
int whA,whB;
int parent[MAX_N];
int sz[MAX_N];
void dfs(int u)
{
    sz[u]=1;
    for(int v:g[u])
    {
        if(v==parent[u])continue;
        deg[u]++;
        deg[v]++;
        parent[v]=u;
        dfs(v);
        sz[u]+=sz[v];
    }
}
void colour(int u,int color)
{
    col[u]=color;
    for(int v:g[u])
    {
        if(v==parent[u])continue;
        colour(v,color);
    }
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q)
{
    vector<int>ans;
    ans.resize(n);

    vector<pair<int,int>>shit;
    shit.push_back({a,1});shit.push_back({b,2});shit.push_back({c,3});
    sort(shit.begin(),shit.end());

    A=shit[0].first;whA=shit[0].second;
    B=shit[1].first;whB=shit[1].second;

    for(int i=0;i<p.size();i++)
    {
        g[p[i]].push_back(q[i]);
        g[q[i]].push_back(p[i]);
    }

    parent[0]=-1;
    dfs(0);

    vector<pair<int,int>>subtrees;
    for(int u:g[0])
    {
        subtrees.push_back({sz[u],u});
    }
    sort(subtrees.begin(),subtrees.end());

    int firstA=-1,firstB=-1;
    for(auto [s,u]:subtrees)
    {
        if(s>=A && firstA==-1)
        {
            firstA=u;
        }
        if(s>=B && firstB==-1)
        {
            firstB=u;
        }
    }

    if(firstA==-1)return ans;

    if(n-sz[firstA]>=B)
    {
        colour(0,whB);
        colour(firstA,whA);
    }
    else if(firstB!=-1 && n-sz[firstB]>=A)
    {
        colour(0,whA);
        colour(firstB,whB);
    }
    else return ans;


    queue<int>bfs;
    int cntwhA=0;
    for(int i=0;i<n;i++)
    {
        if(col[i]==whA)
        {
            cntwhA++;
            if(deg[i]==1)bfs.push(i);
        }
    }
    int toremove=cntwhA-A;
    while(bfs.size() && toremove>0)
    {
        int u=bfs.front();
        bfs.pop();
        toremove--;
        col[u]=shit[2].second;

        for(int v:g[u])
        {
            if(col[v]==whA)
            {
                deg[v]--;
                if(deg[v]==1)bfs.push(v);
            }
        }
    }

    while(bfs.size())bfs.pop();
    int cntwhB=0;
    for(int i=0;i<n;i++)
    {
        if(col[i]==whB)
        {
            cntwhB++;
            if(deg[i]==1)bfs.push(i);
        }
    }
    toremove=cntwhB-B;
    while(bfs.size() && toremove>0)
    {
        int u=bfs.front();
        bfs.pop();
        toremove--;
        col[u]=shit[2].second;

        for(int v:g[u])
        {
            if(col[v]==whB)
            {
                deg[v]--;
                if(deg[v]==1)bfs.push(v);
            }
        }
    }
    //B

    for(int i=0;i<n;i++)ans[i]=col[i];
    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...