Submission #570028

#	Time	Username	Problem	Language	Result	Execution time	Memory
570028		Stickfish	Split the Attractions (IOI19_split)	C++17	40 / 100	1155 ms	16960 KiB

split

#include "split.h"
#include <vector>
#include <random>
#include <cassert>
#include <algorithm>
#include <bitset>
#include <vector>
using namespace std;

const int MAXN = 1e5 + 123;
vector<int> edg[MAXN];
int depth[MAXN];
int sz[MAXN];
int rt[MAXN];
int chld[MAXN];
bitset<MAXN> used;

void dfs(int v) {
    used[v] = 1;
    sz[v] = 1;
    chld[v] = 0;
    for (auto u : edg[v]) {
        if (used[u])
            continue;
        ++chld[v];
        depth[u] = depth[v] + 1;
        rt[u] = v;
        dfs(u);
        sz[v] += sz[u];
    }
}

vector<int> split0(int n, int a, int b) {
    for (int i = 0; i < n; ++i) {
        if (rt[i] == i)
            continue;
        if (sz[i] >= b && n - sz[i] >= a) {
            swap(a, b);
        }
        if (sz[i] >= a && n - sz[i] >= b) {
            vector<int> ans(n, 0);
            int ri = rt[i];
            edg[i].erase(find(edg[i].begin(), edg[i].end(), ri));
            edg[ri].erase(find(edg[ri].begin(), edg[ri].end(), i));
            rt[i] = i;
            rt[ri] = ri;
            for (int i = 0; i < n; ++i)
                used[i] = 0;
            dfs(ri);
            for (int i = 0; i < n; ++i)
                used[i] = 0;
            dfs(i);
            --chld[i];
            --chld[ri];
            vector<int> leafa;
            vector<int> leafb;
            for (int j = 0; j < n; ++j) {
                if (used[j]) {
                    ans[j] = 0;
                    if (chld[j] == 0)
                        leafa.push_back(j);
                    continue;
                }
                ans[j] = 1;
                if (chld[j] == 0) {
                    leafb.push_back(j);
                }
            }
            int c = n - a - b;
            int sza = sz[i];
            int szb = sz[ri];
            while (c > 0 && sza > a && leafa.size()) {
                int v = leafa.back();
                leafa.pop_back();
                --c;
                --sza;
                ans[v] = 2;
                for (auto u : edg[v]) {
                    if (--chld[u] == 0)
                        leafa.push_back(u);
                }
            }
            while (c > 0 && szb > b && leafb.size()) {
                int v = leafb.back();
                leafb.pop_back();
                --c;
                --szb;
                ans[v] = 2;
                for (auto u : edg[v]) {
                    if (--chld[u] == 0)
                        leafb.push_back(u);
                }
            }
            return ans;
        }
    }
    return {};
}

vector<int> solve(int n, int a, int b, int c, int root) {
    for (int i = 0; i < n; ++i)
        used[i] = 0;
    depth[root] = 0;
    rt[root] = root;
    dfs(root);
    vector<int> vls = {a, b, c};
    vector<int> rpl;
    if (c >= a && c >= b)
        rpl = {0, 1, 2};
    else if (b >= a && b >= c)
        rpl = {0, 2, 1};
    else
        rpl = {1, 2, 0};
    if (vls[rpl[0]] > vls[rpl[1]])
        swap(rpl[0], rpl[1]);
    vector<int> res = split0(n, vls[rpl[0]], vls[rpl[1]]);
    if (res.empty()) {
        return {};
    } else {
        vector<int> cnts(3, 0);
        for (int i = 0; i < n; ++i)
            ++cnts[res[i]];
        if (cnts[1] < cnts[0]) {
            for (int i = 0; i < n; ++i) {
                if (res[i] < 2)
                    res[i] = 1 ^ res[i];
            }
        }
        vector<int> rrpl(3);
        for (int i = 0; i < 3; ++i)
            rrpl[rpl[i]] = i;
        for (int i = 0; i < n; ++i)
            res[i] = rpl[res[i]] + 1;
        return res;
    }
	return res;

}

mt19937 rng(177013);

struct dsu {
    void resize(int n) {
        sz.assign(n, 1);
        rt.resize(n);
        for (int i = 0; i < n; ++i)
            rt[i] = i;
    }

    int gst(int i) {
        if (rt[i] == i)
            return i;
        return rt[i] = gst(rt[i]);
    }

    bool unite(int i, int j) {
        i = gst(i);
        j = gst(j);
        if (i == j)
            return false;
        if (sz[i] < sz[j])
            swap(j, i);
        sz[i] += sz[j];
        rt[j] = i;
        return true;
    }
 private:
     vector<int> rt;
     vector<int> sz;
};

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    int t = 10000000 / (n + p.size());
    vector<pair<int, int>> edges;
    for (int i = 0; i < p.size(); ++i)
        edges.push_back({p[i], q[i]});
    while (t--) {
        for (int i = 0; i < n; ++i)
            edg[i].clear();
        shuffle(edges.begin(), edges.end(), rng);
        dsu su;
        su.resize(n);
        for (auto [u, v] : edges) {
            if (su.unite(v, u)) {
                edg[u].push_back(v);
                edg[v].push_back(u);
            }
        }
        int root = rng() % n;
        vector<int> res = solve(n, a, b, c, root);
        if (res.size())
            return res;
    }
    return vector<int>(n);
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:175:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  175 |     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