Submission #265319

#	Time	Username	Problem	Language	Result	Execution time	Memory
265319		hamerin	Split the Attractions (IOI19_split)	C++17	100 / 100	343 ms	35772 KiB

split

#include "split.h"

#include <bits/stdc++.h>

#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

using namespace std;

using i64 = long long;
using d64 = long double;
using pi = pair<int, int>;
using pli = pair<i64, i64>;
using ti = tuple<int, int, int>;
using tli = tuple<i64, i64, i64>;

#define iterall(cont) cont.begin(), cont.end()
#define prec(n) setprecision(n) << fixed

class disjointSet {
   public:
    vector<int> p;

    disjointSet() = default;
    explicit disjointSet(int N) {
        p.clear();
        p.resize(N);
        for (int i = 0; i < N; i++) p[i] = i;
    }

    int find(int u) { return p[u] = (p[u] == u ? u : find(p[u])); }

    void mer(int u, int v) { p[find(v)] = find(u); }

    bool sset(int u, int v) { return find(u) == find(v); }
};

namespace GRAPH {

int V, E, cen;
vector<int> sut, sz, par;
vector<pi> edges;
vector<vector<int>> adj, tr;
vector<bool> vst;
disjointSet dS;

void init(int _V) {
    V = _V;
    sut.resize(V), sz.resize(V), par.resize(V);
    adj.resize(V), tr.resize(V);
    dS = disjointSet(V);
}

void emplace_edge(int u, int v) {
    ++E;
    edges.emplace_back(u, v);
    adj[u].emplace_back(v);
    adj[v].emplace_back(u);
}

// dfs ordering && parent
void dfs1(int h, bool ini = false) {
    if (ini) {
        vst.clear();
        vst.resize(V);
        vst[h] = true;
        par[h] = -1;
    }

    for (auto t : adj[h]) {
        if (vst[t]) continue;

        vst[t] = true;
        tr[h].emplace_back(t);
        par[t] = h;
        dfs1(t);
    }
}

// subtree size
void dfs2(int h) {
    sz[h] = 1;
    for (auto t : tr[h]) {
        dfs2(t);
        sz[h] += sz[t];
    }
}

// find centroid
int dfs3(int h) {
    for (auto t : tr[h])
        if (sz[t] > (V >> 1)) return dfs3(t);
    return h;
}

bool trAdj(int u, int v) { return par[u] == v || par[v] == u; }

// centroid - dfs coloring
void dfs4(int h, int p) {
    for (auto t : adj[h]) {
        if (!trAdj(h, t)) continue;
        if (t == p) continue;

        if (h != cen) dS.mer(h, t);
        dfs4(t, h);
    }
}

void dfs() {
    dfs1(0, true);
    dfs2(0);
    cen = dfs3(0);
    dfs4(cen, -1);

    // cout << cen << endl;
    for (int i = 0; i < V; i++) sut[i] = dS.find(i);
}

void con(int h, int p, vector<int> &r, int &n) {
    if (n == r.size()) return;
    r[n++] = h;

    for (auto t : adj[h]) {
        if (!trAdj(h, t)) continue;
        if (t == p) continue;

        con(t, h, r, n);
    }
}
}  // namespace GRAPH

vector<int> find_split(int n, int a, int b, int c, vector<int> p,
                       vector<int> q) {
    vector<pi> targetSize = {pi(a, 1), pi(b, 2), pi(c, 3)};
    sort(iterall(targetSize));

    const size_t M = p.size();
    GRAPH::init(n);
    for (int i = 0; i < M; i++) GRAPH::emplace_edge(p[i], q[i]);
    GRAPH::dfs();

    auto newVer = GRAPH::sut;
    sort(iterall(newVer));

    vector<int> vertexSize(n);
    for (auto el : newVer) vertexSize[el]++;

    newVer.erase(unique(iterall(newVer)), newVer.end());
    int W = newVer.size();

    vector<int> aConf;

    // a이상 서브트리 존재
    for (int i = 0; i < W; i++) {
        if (newVer[i] == GRAPH::cen) continue;
        if (vertexSize[newVer[i]] >= targetSize[0].first) {
            aConf.emplace_back(newVer[i]);
            break;
        }
    }

    // a이상 서브트리 존재X
    if (aConf.empty()) {
        auto dS = disjointSet(W);
        vector<set<int>> newAdj(W);

        for (auto [u, v] : GRAPH::edges) {
            if (GRAPH::trAdj(u, v) || u == GRAPH::cen || v == GRAPH::cen)
                continue;

            auto cu =
                lower_bound(iterall(newVer), GRAPH::sut[u]) - newVer.begin();
            auto cv =
                lower_bound(iterall(newVer), GRAPH::sut[v]) - newVer.begin();
            dS.mer(cu, cv);
            if (cu != cv) newAdj[cu].emplace(cv), newAdj[cv].emplace(cu);
        }

        vector<int> dSroot(W);
        for (int i = 0; i < W; i++) {
            if (newVer[i] == GRAPH::cen) continue;
            dSroot[i] = dS.find(i);
        }

        vector<int> sz(W);
        for (int i = 0; i < W; i++) {
            if (newVer[i] == GRAPH::cen) continue;
            sz[dSroot[i]] += vertexSize[newVer[i]];
        }

        for (int i = 0; i < W; i++) {
            if (sz[i] >= targetSize[0].first) {
                int vs = 0;

                {
                    queue<int> que;
                    vector<bool> vst(W, true);
                    for (int j = 0; j < W; j++) if (dSroot[j] == i) vst[j] = false;

                    vst[i] = true;
                    que.emplace(i);

                    while (!que.empty() && vs < targetSize[0].first) {
                        int cur = que.front();
                        que.pop();

                        aConf.emplace_back(newVer[cur]);
						vs += vertexSize[newVer[cur]];

						for (auto there : newAdj[cur]) {
                            // cout << there << endl;
                            if (vst[there]) continue;
                            que.emplace(there);
                            vst[there] = true;
                        }
                    }
                }

                break;
            }
        }
    }

    // Configure
    if (aConf.empty()) return vector<int>(n, 0);

    vector<int> av(n, -1);
    int an = 0;
    for (auto el : aConf) GRAPH::con(el, GRAPH::cen, av, an);

    while (av.back() == -1) av.pop_back();

    // cout << aConf.size() << endl;

    vector<int> avv;
    {
        queue<int> que;
        vector<bool> vst(n, true);
        for (auto el : av) vst[el] = false;

        vst[av[0]] = true;
        que.emplace(av[0]);

        while (!que.empty() && avv.size() < targetSize[0].first) {
            int cur = que.front();
            que.pop();

            avv.emplace_back(cur);
            for (auto there : GRAPH::adj[cur]) {
                // cout << there << endl;
                if (vst[there]) continue;
                que.emplace(there);
                vst[there] = true;
            }
        }
    }

    vector<int> bvv;
    {
        queue<int> que;
        vector<bool> vst(n);
        for (auto el : avv) vst[el] = true;

        vst[GRAPH::cen] = true;
        que.emplace(GRAPH::cen);

        while (!que.empty() && bvv.size() < targetSize[1].first) {
            int cur = que.front();
            que.pop();

            bvv.emplace_back(cur);
            for (auto there : GRAPH::adj[cur]) {
                if (vst[there]) continue;
                que.emplace(there);
                vst[there] = true;
            }
        }
    }

    vector<int> ret(n, targetSize[2].second);
    for (auto el : avv) ret[el] = targetSize[0].second;
    for (auto el : bvv) ret[el] = targetSize[1].second;

    return ret;
}

Compilation message (stderr)

split.cpp: In function 'void GRAPH::con(int, int, std::vector<int>&, int&)':
split.cpp:120:11: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  120 |     if (n == r.size()) return;
      |         ~~^~~~~~~~~~~
split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:139:23: warning: comparison of integer expressions of different signedness: 'int' and 'const size_t' {aka 'const long unsigned int'} [-Wsign-compare]
  139 |     for (int i = 0; i < M; i++) GRAPH::emplace_edge(p[i], q[i]);
      |                     ~~^~~
split.cpp:244:43: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  244 |         while (!que.empty() && avv.size() < targetSize[0].first) {
      |                                ~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~
split.cpp:267:43: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  267 |         while (!que.empty() && bvv.size() < targetSize[1].first) {
      |                                ~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~

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