Submission #826755

#	Time	Username	Problem	Language	Result	Execution time	Memory
826755		kevinsogo	Split the Attractions (IOI19_split)	C++17	100 / 100	274 ms	77544 KiB

split

#include <bits/stdc++.h>
using namespace std;
#include "split.h"
using ll = long long;
template<class T>
using vec = vector<T>;

vec<vec<int>> extract(int n, int a, int b, int c, const vec<pair<int,int>>& edges) {
    int e = edges.size();
    assert(1 <= a and a <= b and b <= c);
    assert(3*a <= n);
    assert(2*b <= n);
    assert(1*c <= n);
    
    vec<vec<int>> adj(n);
    vec<bool> vis(n);
    auto get_component = [&](unordered_set<int>& nodes, int s) {
        for (int i : nodes) {
            vis[i] = false;
            adj[i].clear();
        }
        assert(nodes.count(s));
        for (auto [i, j] : edges) {
            if (nodes.count(i) and nodes.count(j)) {
                adj[i].push_back(j);
                adj[j].push_back(i);
            }
        }

        vis[s] = true;
        vec<int> que = {s};
        for (int f = 0; f < que.size(); f++) {
            int i = que[f];
            for (int j : adj[i]) {
                if (not vis[j]) {
                    vis[j] = true;
                    que.push_back(j);
                }
            }
        }
        return que;
    };

    auto get_bccs = [&]() {
        vec<vec<pair<int,int>>> adj(n);
        for (int idx = 0; idx < e; idx++) {
            auto [i, j] = edges[idx];
            adj[i].emplace_back(j, idx);
            adj[j].emplace_back(i, idx);
        }

        vec<int> low(n, -1), dsc(n, -1);
        vec<bool> evi(e);
        vec<int> eve;
        vec<vec<int>> bccs;
        int tm = 0;
        function<void(int)> dfs = [&](int i) {
            assert(dsc[i] == -1);
            dsc[i] = low[i] = tm++;
            for (auto [j, idx] : adj[i]) {
                if (evi[idx]) continue;
                eve.push_back(idx);
                evi[idx] = true;
                if (dsc[j] == -1) {
                    dfs(j);
                    low[i] = min(low[i], low[j]);
                    if (low[j] >= dsc[i]) {
                        bccs.emplace_back();
                        while (bccs.back().empty() or bccs.back().back() != idx) {
                            bccs.back().push_back(eve.back()); eve.pop_back();
                        }
                    }
                } else {
                    low[i] = min(low[i], dsc[j]);
                }
            }
        };

        dfs(0);

        assert(eve.empty());
        // assert all(d >= 0 for d in dsc)
        // assert all(evi)
        // assert not eve
        // assert sum(map(len, bccs)) == e

        return bccs;
    };

    auto bccs = get_bccs();

    auto get_bctree = [&]() -> tuple<vec<vec<int>>,vec<int>,vec<int>> {
        vec<vec<int>> bcadj(n + bccs.size());
        vec<unordered_set<int>> bcedges(n);
        for (int bci = 0; bci < bccs.size(); bci++) {
            for (int idx : bccs[bci]) {
                auto [i, j] = edges[idx];
                bcedges[i].insert(n + bci);
                bcedges[j].insert(n + bci);
            }
        }

        int ct = 0;
        for (int i = 0; i < n; i++) {
            for (int j : bcedges[i]) {
                ct++;
                bcadj[i].push_back(j);
                bcadj[j].push_back(i);
            }
        }
        assert(ct == n + bccs.size() - 1);

        vec<int> bcsiz(n, 1);
        bcsiz.resize(n + bccs.size());
        vec<int> bcpar(n + bccs.size(), -1);
        function<void(int)> dfs = [&](int i) {
            for (int j : bcadj[i]) {
                if (j != bcpar[i]) {
                    bcpar[j] = i;
                    dfs(j);
                    bcsiz[i] += bcsiz[j];
                }
            }
        };

        dfs(0);
        assert(bcsiz[0] == n);
        return {bcadj, bcsiz, bcpar};
    };

    auto [bcadj, bcsiz, bcpar] = get_bctree();

    vec<int> low(n, -1), dsc(n, -1);
    vec<bool> evi(e);
    vec<vec<pair<int,int>>> iadj(n);
    vec<int> isize(n), dep(n);
    vec<bool> inseq(n), inaseq(n);
    vec<int> pare(n, -1);
    vec<vec<int>> igs(n);
    for (int bci = 0; bci < bccs.size(); bci++) {
        vec<int>& inodes = bcadj[n + bci];
        for (int i : inodes) isize[i] = bcsiz[i];
        isize[bcpar[n + bci]] = n - bcsiz[n + bci];
        // assert sum(isize[i] for i in inodes) == n
        assert(inodes.size() >= 2);

        int r = *max_element(inodes.begin(), inodes.end(), [&](int i, int j) { return isize[i] < isize[j]; });

        for (int i : inodes) {
            dsc[i] = -1;
            iadj[i].clear();
            inseq[i] = inaseq[i] = false;
        }
        
        for (int idx : bccs[bci]) {
            evi[idx] = false;
            auto [i, j] = edges[idx];
            iadj[i].emplace_back(j, idx);
            iadj[j].emplace_back(i, idx);
        }
        
        int tm = 0;
        dep[r] = 0;
        function<void(int)> dfs = [&](int i) {
            assert(dsc[i] == -1);
            dsc[i] = low[i] = tm++;
            for (auto [j, idx] : iadj[i]) {
                if (evi[idx]) continue;
                evi[idx] = true;
                if (dsc[j] == -1) {
                    dep[j] = dep[i] + 1;
                    pare[j] = idx;
                    dfs(j);
                    low[i] = min(low[i], low[j]);
                } else {
                    low[i] = min(low[i], dsc[j]);
                }
            }

            assert((low[i] == dsc[i]) == (i == r or inodes.size() == 2));
        };

        dfs(r);

        vec<int> aseq;
        if (inodes.size() == 2) {
            aseq = inodes;
            if (aseq[0] != r) reverse(aseq.begin(), aseq.end());
        } else {
            vec<int> fseq;
            auto add_all = [&](vec<int> s) {
                reverse(s.begin(), s.end());
                for(int i : s) {
                    assert(not inseq[i]);
                    inseq[i] = true;
                }
                for (int i : s) {
                    igs[i].clear();
                    for (auto [j, idx] : iadj[i]) {
                        if (not inseq[j] and pare[j] == idx) igs[i].push_back(j);
                    }
                    fseq.push_back(i);
                }
            };
            add_all({r});
            while (not fseq.empty()) {
                int i = fseq.back(); fseq.pop_back();
                if (igs[i].empty()) {
                    aseq.push_back(i);
                    inaseq[i] = true;
                } else {
                    int j = igs[i].back(); igs[i].pop_back();
                    int targ = max(1, dsc[i]);
                    assert(low[j] < targ);
                    vec<int> seq = {j};
                    while (dsc[seq.back()] >= targ) {
                        for (auto [k, idx] : iadj[seq.back()]) {
                            if (dsc[k] < targ or pare[k] == idx and low[k] == low[seq.back()]) {
                                seq.push_back(k);
                                break;
                            }
                        }
                        // else: assert false
                    }
                    int l = seq.back(); seq.pop_back();
                    if (inaseq[l]) {
                        reverse(seq.begin(), seq.end());
                        fseq.push_back(i);
                        add_all(seq);
                    } else {
                        add_all(seq);
                        fseq.push_back(i);
                    }
                }
            }
        }

        // assert {*aseq} == {*inodes}

        int sz = 0;
        vec<int> gr, rg(aseq.rbegin(), aseq.rend());
        while (sz < a) {
            gr.push_back(rg.back()); rg.pop_back();
            sz += isize[gr.back()];
        }

        assert(not gr.empty());
        assert(not rg.empty());

        if (b <= sz and a <= n - sz) {
            swap(gr, rg);
            sz = n - sz;
            assert(a <= sz and b <= n - sz);
        }

        if (a <= sz and b <= n - sz) {
            unordered_set<int> ns;
            for (int i = 0; i < n; i++) ns.insert(i);
            for (int i : gr) ns.erase(i);
            auto compb = get_component(ns, rg.front());
            assert(compb.size() == n - sz);
            ns.clear();
            for (int i = 0; i < n; i++) ns.insert(i);
            for (int i : compb) ns.erase(i);
            auto compa = get_component(ns, gr.front());
            assert(compa.size() == sz);
            compa.resize(a);
            compb.resize(b);
            ns.clear();
            for (int i = 0; i < n; i++) ns.insert(i);
            for (int i : compa) ns.erase(i);
            for (int i : compb) ns.erase(i);
            vec<int> compc(ns.begin(), ns.end());
            return {compa, compb, compc};
        }
    }

    return {};
}

vec<int> find_split(int n, int a, int b, int c, vec<int> p, vec<int> q) {
    assert(p.size() == q.size());
    vec<pair<int,int>> edges(p.size());
    for (int i = 0; i < p.size(); i++) edges[i] = {p[i], q[i]};
    vec<pair<int,int>> abci = {{a, 1}, {b, 2}, {c, 3}};
    sort(abci.begin(), abci.end());
    vec<int> res(n, 0);
    if (auto abcs = extract(n, abci[0].first, abci[1].first, abci[2].first, edges); not abcs.empty()) {
        assert(abci.size() == abcs.size());
        for (int x = 0; x < abci.size(); x++) {
            auto [sz, vl] = abci[x];
            assert(abcs[x].size() == sz);
            for (int i : abcs[x]) res[i] = vl;
        }
    }
    return res;
}

Compilation message (stderr)

split.cpp: In lambda function:
split.cpp:32:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   32 |         for (int f = 0; f < que.size(); f++) {
      |                         ~~^~~~~~~~~~~~
split.cpp: In lambda function:
split.cpp:95:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int>, std::allocator<std::vector<int> > >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   95 |         for (int bci = 0; bci < bccs.size(); bci++) {
      |                           ~~~~^~~~~~~~~~~~~
In file included from /usr/include/c++/10/cassert:44,
                 from /usr/include/x86_64-linux-gnu/c++/10/bits/stdc++.h:33,
                 from split.cpp:1:
split.cpp:111:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int>, std::allocator<std::vector<int> > >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  111 |         assert(ct == n + bccs.size() - 1);
      |                ~~~^~~~~~~~~~~~~~~~~~~~~~
split.cpp: In function 'vec<std::vector<int> > extract(int, int, int, int, vec<std::pair<int, int> >&)':
split.cpp:140:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int>, std::allocator<std::vector<int> > >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  140 |     for (int bci = 0; bci < bccs.size(); bci++) {
      |                       ~~~~^~~~~~~~~~~~~
split.cpp:218:65: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  218 |                             if (dsc[k] < targ or pare[k] == idx and low[k] == low[seq.back()]) {
In file included from /usr/include/c++/10/cassert:44,
                 from /usr/include/x86_64-linux-gnu/c++/10/bits/stdc++.h:33,
                 from split.cpp:1:
split.cpp:261:33: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  261 |             assert(compb.size() == n - sz);
      |                    ~~~~~~~~~~~~~^~~~~~~~~
split.cpp:266:33: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  266 |             assert(compa.size() == sz);
      |                    ~~~~~~~~~~~~~^~~~~
split.cpp: In function 'vec<int> find_split(int, int, int, int, vec<int>, vec<int>)':
split.cpp:284:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  284 |     for (int i = 0; i < p.size(); i++) edges[i] = {p[i], q[i]};
      |                     ~~^~~~~~~~~~
split.cpp:290:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int>, std::allocator<std::pair<int, int> > >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  290 |         for (int x = 0; x < abci.size(); x++) {
      |                         ~~^~~~~~~~~~~~~
In file included from /usr/include/c++/10/cassert:44,
                 from /usr/include/x86_64-linux-gnu/c++/10/bits/stdc++.h:33,
                 from split.cpp:1:
split.cpp:292:35: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'std::tuple_element<0, std::pair<int, int> >::type' {aka 'int'} [-Wsign-compare]
  292 |             assert(abcs[x].size() == sz);
      |                    ~~~~~~~~~~~~~~~^~~~~

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