Submission #570137

#	Time	Username	Problem	Language	Result	Execution time	Memory
570137		Stickfish	Split the Attractions (IOI19_split)	C++17	100 / 100	194 ms	26956 KiB

split

#include "split.h"
#include <vector>
#include <iostream>
#include <cassert>
#include <algorithm>
#include <bitset>
#include <vector>
using namespace std;
 
const int MAXN = 1e5 + 123;
vector<int> edg[MAXN];
vector<int> edg_st[MAXN];
int sz[MAXN];
int rt[MAXN];
int timer = 0;
int tin[MAXN];
int fup[MAXN];
bitset<MAXN> used;
 
void dfs(int v) {
    fup[v] = tin[v] = timer++;
    used[v] = 1;
    sz[v] = 1;
    for (auto u : edg[v]) {
        if (u == rt[v]) {
            continue;
        }
        if (used[u]) {
            fup[v] = min(fup[v], tin[u]);
            continue;
        }
        edg_st[v].push_back(u);
        rt[u] = v;
        dfs(u);
        sz[v] += sz[u];
        fup[v] = min(fup[v], fup[u]);
    }
}

int get_centr(int v, int n) {
    for (auto u : edg_st[v]) {
        if (sz[u] > n / 2)
            return get_centr(u, n);
    }
    return v;
}

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> strip_all(int n, int a, int b, vector<int> res) {
    dsu su;
    su.resize(n);
    for (int i = 0; i < n; ++i)
        edg_st[i].clear();
    for (int v = 0; v < n; ++v) {
        for (auto u : edg[v]) {
            if (res[u] == res[v]) {
                if (su.unite(u, v)) {
                    edg_st[u].push_back(v);
                    edg_st[v].push_back(u);
                }
            }
        }
    }
    vector<int> dg(n, 0);
    for (int v = 0; v < n; ++v) {
        for (auto u : edg_st[v]) {
            if (res[v] == res[u]) {
                ++dg[v];
            }
        }
        //cout << dg[v] << ' ';
    }
    //cout << endl;
    vector<int> leafa;
    vector<int> leafb;
    int sza = 0;
    int szb = 0;
    for (int j = 0; j < n; ++j) {
        if (res[j] == 0) {
            ++sza;
            if (dg[j] == 1)
                leafa.push_back(j);
        } else {
            ++szb;
            if (dg[j] == 1)
                leafb.push_back(j);
        }
    }
    assert(sza <= szb);
    assert(sza >= a && szb >= b);
    while (sza > a) {
        int v = leafa.back();
        leafa.pop_back();
        --sza;
        for (auto u : edg_st[v]) {
            if (res[u] == res[v] && --dg[u] == 1)
                leafa.push_back(u);
        }
        res[v] = 2;
    }
    while (szb > b) {
        int v = leafb.back();
        leafb.pop_back();
        --szb;
        for (auto u : edg_st[v]) {
            if (res[u] == res[v] && --dg[u] == 1)
                leafb.push_back(u);
        }
        res[v] = 2;
    }
    return res;
}

void recolor_dfs(int v, int cl, vector<int>& res) {
    res[v] = cl;
    for (auto u : edg_st[v])
        recolor_dfs(u, cl, res);
}

 
vector<int> split0(int n, int a, int b) {
    //for (int i = 0; i < n; ++i) {
        //for (auto u : edg[i])
            //cout << u << ' ';
        //cout << endl;
    //}
    for (int i = 0; i < n; ++i) {
        if (rt[i] == i)
            continue;
        if (sz[i] >= a && sz[i] <= n - a) {
            vector<int> ans(n, 1);
            recolor_dfs(i, 0, ans);
            if (sz[i] > n / 2) {
                for (int j = 0; j < n; ++j)
                    ans[j] ^= 1;
            }
            for (int i = 1; i < n; ++i)
                edg_st[i].push_back(rt[i]);
            //for (int i = 0; i < n; ++i)
                //cout << ans[i] << ' ';
            //cout << endl;
            return strip_all(n, a, b, ans);
        }
    }
    int ctr = get_centr(0, n);
    int mnsz = 1;
    vector<int> res(n, 1);
    res[ctr] = 0;
    for (auto u : edg_st[ctr]) {
        if (fup[u] >= tin[ctr]) {
            mnsz += sz[u];
            recolor_dfs(u, 0, res);
        }
    }
    if (mnsz > n - a)
        return {};
    for (auto u : edg_st[ctr]) {
        if (fup[u] < tin[ctr] && mnsz < a) {
            mnsz += sz[u];
            recolor_dfs(u, 0, res);
        }
    }
    assert(mnsz <= n - a);
    if (mnsz > n / 2) {
        for (int i = 0; i < n; ++i)
            res[i] ^= 1;
        mnsz = n - mnsz;
    }
    return strip_all(n, a, b, res);
}
 
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    for (int i = 0; i < n; ++i)
        edg[i].clear();
    for (int i = 0; i < p.size(); ++i) {
        edg[p[i]].push_back(q[i]);
        edg[q[i]].push_back(p[i]);
    }
    used = 0;
    dfs(0);
    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]);
    assert(vls[rpl[0]] <= vls[rpl[1]] && vls[rpl[1]] <= vls[rpl[2]]);
    vector<int> res = split0(n, vls[rpl[0]], vls[rpl[1]]);
    if (res.empty()) {
        res.assign(n, 0);
    } else {
        vector<int> cnts(3, 0);
        for (int i = 0; i < n; ++i)
            ++cnts[res[i]];
        assert(cnts[2] >= cnts[1] && cnts[2] >= cnts[0]);
        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;
}

Compilation message (stderr)

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

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