Submission #1049105

#	Time	Username	Problem	Language	Result	Execution time	Memory
1049105		dead0ne	Split the Attractions (IOI19_split)	C++17	40 / 100	66 ms	22732 KiB

split

#pragma GCC optimize("unroll-loops,Ofast,O3")
#include <bits/stdc++.h>
#define pb push_back
#define mp make_pair
#define spc << " " <<
#define endl "\n"
#define all(x) x.begin(), x.end()
//#define int long long
#define ii pair<long long,int>
#define vi vector<int>
#define vii vector<ii>
#define st first
#define nd second
#define mid (l+r)/2
#define inf 1e15
#define MOD 1000000007
#define MX 100005
using namespace std;


int tim, baddy, cur;
vii edges[MX], colors(3);
int vis[MX], sub[MX], in[MX], low[MX], good[MX];
vi ans, colored(3);
void tree(int node, int par){
    vis[node]=sub[node]=1;
    tim++;
    in[node]=low[node]=tim;
    //cerr << "got into" spc node spc in[node] << endl;
    for(auto &i:edges[node]){
        if(vis[i.st]){
            if(i.st!=par){
                //cerr << "back edge" spc node spc i.st << endl;
                low[node] = min(low[node], in[i.st]);
            }
            else{
                i.nd = 1;
            }
            continue;
        }
        i.nd = 1;
        tree(i.st, node);
        sub[node] += sub[i.st];
        //low[node] = min(low[node], low[i.st]);
    }
    //cerr << "getting out of" spc node spc low[node] << endl;
}
void paint_below(int node, int color){
    if(colored[color] >= colors[color].st || ans[node]!=-1) return;
    ans[node]=color;
    colored[color]++;
    for(auto i:edges[node]){
        if(i.nd==0 || in[i.st]<in[node] || ans[i.st]!=-1) continue;
        paint_below(i.st, color);
    }
}
void finalzz(int node){
    if(cur >= colors[0].st) return;
    if(good[node]==1 && low[node]<in[baddy]){
        cur+=sub[node];
        paint_below(node, 0);
        return;
    } 
    for(auto i:edges[node]){
        if(i.nd==0 || in[i.st]<in[node]) continue;
        finalzz(i.st);
    }
}

vi find_split(int n, int a, int b, int c, vi p, vi q){
    colors = {{a,1}, {b,2}, {c,3}};
    sort(all(colors));
    colored = {0,0,0};

    ans.resize(n,-1);
    for(int i=0; i<n; i++){
        edges[i].clear();
        vis[i]=sub[i]=0;
        good[i]=-1;
    }
    for(int i=0; i<p.size(); i++){
        edges[p[i]].pb({q[i],0});
        edges[q[i]].pb({p[i],0});
    }

    tim=0;
    tree(0, 0);
    /*for(int i=0; i<n; i++){
        //cerr << i spc sub[i] spc in[i] spc low[i] << endl;
    }*/

    for(int i=1; i<n; i++){
        if(sub[i]>=colors[0].st && sub[i]<=colors[0].st+colors[2].st){
            //cerr << "subtree fitting a" spc i << endl;
            paint_below(i, 0);
            paint_below(0, 1);
            for(int j=0; j<n; j++) if(ans[j]==-1) ans[j]=2;
            for(int k=0; k<n; k++) ans[k] = colors[ans[k]].nd;
            return ans;
        }
        else if(sub[i]>=colors[1].st && sub[i]<=colors[1].st+colors[2].st){
            //cerr << "subtree fitting b" spc i << endl;
            paint_below(i, 1);
            paint_below(0, 0);
            for(int j=0; j<n; j++) if(ans[j]==-1) ans[j]=2;
            for(int k=0; k<n; k++) ans[k] = colors[ans[k]].nd;
            return ans;
        }
        else if(sub[i]<colors[0].st) good[i]=1;
        else if(sub[i]>colors[1].st+colors[2].st) good[i]=0;
        else assert(false);
    }
    baddy=0;
    for(int i=1; i<n; i++) if(good[i]==0 && in[i]>in[baddy]) baddy=i; //are we the baddies???
    int cur = n-sub[baddy];
    assert(cur < colors[0].st);

    finalzz(baddy);
    paint_below(baddy, 1);
    paint_below(0, 0);
    if(colored[0]!=colors[0].st || colored[1]!=colors[1].st){
        fill(ans.begin(), ans.end(), 0);
        return ans;
    }
    for(int j=0; j<n; j++) if(ans[j]==-1) ans[j]=2;
    for(int k=0; k<n; k++) ans[k] = colors[ans[k]].nd;
    return ans;
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:81:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   81 |     for(int i=0; i<p.size(); i++){
      |                  ~^~~~~~~~~

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