Submission #260281

#TimeUsernameProblemLanguageResultExecution timeMemory
260281anonymousSplit the Attractions (IOI19_split)C++14
100 / 100
185 ms26356 KiB
#include "split.h"
#include <vector>
#include <utility>
#include <algorithm>
#define MAXN 100005
using namespace std;

int N, M, CanDo = 1, Done,  Ans[MAXN], NodeType[MAXN];
int Size[MAXN], Lo[MAXN], Dep[MAXN], Par[MAXN];
pair <int,int> Days[3];
vector <int> adj[MAXN];
vector <pair<int,bool> > spT[MAXN]; //spanning tree rooted at 1

void dfs2(int u, int day, bool b) { //b = 1 means dfs on spanning tree
    if (!Days[day].first || Ans[u]) {return;}
    Ans[u] = Days[day].second;
    Days[day].first--;
    if (b) {
        for (auto e: spT[u]) {
            dfs2(e.first, day, b);
        }
    } else {
        for (int v: adj[u]) {
            dfs2(v, day, b);
        }
    }
}

void dfs(int u) {
    if (Done) {return;}
    Size[u] = 1;
    Lo[u] = u;
    bool AllSmall = 1;
    int cutsize = 1;
    for (int v: adj[u]) {
        if (!Size[v]) {
            Dep[v] = Dep[u] + 1;
            Par[v] = u;
            dfs(v);
            if (Dep[Lo[v]] < Dep[Lo[u]]) {
                Lo[u] = Lo[v];
            }
            AllSmall &= Size[v] < Days[0].first;
            if (Lo[v] == v || Lo[v] == u) {
                spT[u].push_back({v,1});
                cutsize += Size[v];
            } else {
                spT[u].push_back({v,0});
            }
            Size[u] += Size[v];
        } else if (v != Par[u] && Dep[v] < Dep[Lo[u]]) {
            Lo[u] = v;
        }
    }
    if (Size[u] >= Days[0].first && Size[u] <= N-Days[0].first) {
        if (Size[u] >= Days[1].first) {
            dfs2(u, 1, 1);
            dfs2(Par[u], 0, 0);
        } else {
            dfs2(u,0, 1);
            dfs2(Par[u],1, 0);
        }
        Done = true;
    } else if ((Size[u] > N-Days[0].first) && AllSmall) {
        if (cutsize > N-Days[0].first) {
            //impossible
            CanDo = 0;
        } else {
            //we got this
            if (cutsize >= Days[1].first) {
                Ans[u] = Days[1].second;
                Days[1].first--;
                for (auto e: spT[u]) {
                    if (e.second) {
                        dfs2(e.first, 1, 1);
                    }
                }
                dfs2(Par[u], 0, 0);
            } else {
                Ans[u] = Days[0].second;
                Days[0].first--;
                for (auto e: spT[u]) {
                    if (e.second) {
                        dfs2(e.first, 0, 1);
                    }
                }
                for (auto e: spT[u]) {
                    if (!e.second) {
                        dfs2(e.first, 0, 1);
                    }
                }
                dfs2(Par[u], 1, 0);
            }
        }
        Done = true;
    }
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	N = n, M = p.size(), Days[0].first = a, Days[1].first = b, Days[2].first = c;
	for (int i=0; i<3; i++) {Days[i].second = i+1;}
	sort(Days,Days+3);

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

    dfs(0);

    vector <int> Out;
    if (CanDo) {
        for (int i=0; i < N; i++) {
            if (Ans[i]) {
                Out.push_back(Ans[i]);
            } else {
                Out.push_back(Days[2].second);
            }
        }
    } else {
        for (int i=0; i < N; i++) {
            Out.push_back(0);
        }
    }
    return(Out);
}
#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...