Submission #1068918

#	Time	Username	Problem	Language	Result	Execution time	Memory
1068918		RaresFelix	Split the Attractions (IOI19_split)	C++17	100 / 100	156 ms	33572 KiB

split

#include "split.h"
#include <bits/stdc++.h>

using namespace std;

using vi = vector<int>;
using ii = pair<int, int>;

struct DSU {
    vi e;
    DSU(int n) : e(n, -1) {}

    int repr(int u) {
        while(e[u] >= 0) u = e[u];
        return u;
    }

    bool join(int u, int v) {
        u = repr(u);
        v = repr(v);
        if(u == v) return false;
        if(e[u] >= e[v]) swap(u, v);
        e[u] += e[v];
        e[v] = u;
        return true;
    }
};

struct tree {
    int n;
    vi sz;
    vector<vi> L;

    tree(int N) : n(N), L(N), sz(N, 0) {}

    void add_edge(int u, int v) {
        L[u].push_back(v);
        L[v].push_back(u);
    }

    void get_cent(vector<vi> &seturi, vi &Cul, int &cent) {
        ///presupune ca lucram de fapt cu un arbore
        
        function<void(int, int)> dfs_sz = [&](int u, int p) {
            sz[u] = 1;
            for(auto it : L[u]) {
                if(it == p) continue;
                dfs_sz(it, u);
                sz[u] += sz[it];
            }
        };

        function<int(int, int, int)> find_c = [&](int u, int p, int vlim) {
            for(auto it : L[u]) {
                if(it == p) continue;
                if(sz[it] * 2 > vlim) return find_c(it, u, vlim);
            }
            return u;
        };

        function<void(int, int, int)> dfs_cul = [&](int u, int p, int cul) {
            seturi[cul].push_back(u);
            Cul[u] = cul;
            for(auto it : L[u]) {
                if(it == p) continue;
                dfs_cul(it, u, cul);
            }
        };

        dfs_sz(0, -1);

        cent = find_c(0, -1, sz[0]);

        Cul.assign(n, 0);

        int nr_cul = 0;
        for(auto it : L[cent]) {
            seturi.emplace_back();
            dfs_cul(it, cent, nr_cul);
            ++nr_cul;
        }
    }

    vi select(vi weight, vi activ, int lim) {
        vi sol;
        for(int i = 0; i < n; ++i) {
            if(activ[i] && weight[i] >= lim) return vi{i};
        }
        vi viz(n, 0);

        vi S;
        int sum = 0, ok = 0;
        function<void(int, int)> dfs = [&](int u, int p) {
            viz[u] = 1;
            S.push_back(u);
            sum += weight[u];
            if(sum >= lim) {
                ok = 1;
                return;
            }
            for(auto it : L[u]) {
                if(it == p) continue;
                if(activ[it])
                    dfs(it, u);
                if(ok) return;
            }
        };

        for(int i = 0; i < n; ++i) {
            if(!viz[i] && activ[i]) {
                dfs(i, -1);
                if(ok) return S;
                sum = 0;
                S.clear();
            }
        }

        return vi();
    }
};

vi find_split(int n, int a, int b, int c, vi p, vi q) {
    vector<ii> Mar = {{a, 1}, {b, 2}, {c, 3}};
    sort(Mar.rbegin(), Mar.rend());

    int m = int(p.size());
    tree T(n);
    DSU helper(n);
    for(int i = 0; i < m; ++i) {
        if(helper.join(p[i], q[i])) {
            T.add_edge(p[i], q[i]);
        }
    }

    int cent;
    vi Cul;
    vector<vi> Multimi;

    T.get_cent(Multimi, Cul, cent);
    int nr_cul = int(Multimi.size());

    DSU UF_cul(nr_cul);
    tree T_cul(nr_cul);
    Cul[cent] = -1;
    for(int i = 0; i < m; ++i) {
        int c1 = Cul[p[i]], c2 = Cul[q[i]];
        if(c1 == -1 || c2 == -1) continue; // cu cent
        if(UF_cul.join(c1, c2)) T_cul.add_edge(c1, c2);
    }

    vi Weight(nr_cul);
    for(int i = 0; i < nr_cul; ++i) {
        Weight[i] = (int)Multimi[i].size();
    }

    auto Cul_C = T_cul.select(Weight, vi(nr_cul, 1), Mar[2].first);
    if(Cul_C.empty()) return vi(n, 0);

    set<int> NoduriC, NoduriB;
    for(auto it : Cul_C) {
        copy(Multimi[it].begin(), Multimi[it].end(), inserter(NoduriC, NoduriC.begin()));
    }
    for(int i = 0; i < n; ++i)
        if(!NoduriC.count(i)) NoduriB.insert(i);

    auto restrict = [&](set<int> &S, int lim) {
        vi active(n, 0), weight(n, 1);
        for(auto it : S) active[it] = 1;
        tree Tnou(n);
        DSU Hnou(n);
        for(int i = 0; i < m; ++i) {
            int u = p[i], v = q[i];
            if(!active[u] || !active[v]) continue;
            if(Hnou.join(u, v)) {
                Tnou.add_edge(u, v);
            }
        }

        auto alese = Tnou.select(weight, active, lim);
        S.clear();
        for(auto it : alese) S.insert(it);
    };

    restrict(NoduriB, Mar[1].first);
    restrict(NoduriC, Mar[2].first);

    vi Re(n, Mar[0].second);
    for(auto it : NoduriB) Re[it] = Mar[1].second;
    for(auto it : NoduriC) Re[it] = Mar[2].second;

   // for(auto it : NoduriB)
   //     cout << it << " ";
   // cout << "\n";
   // for(auto it : NoduriC)
   //     cout << it << " ";
   // cout << "\n";

    assert(NoduriB.size() == Mar[1].first);
    assert(NoduriC.size() == Mar[2].first);
    
    return Re;
}

Compilation message (stderr)

split.cpp: In constructor 'tree::tree(int)':
split.cpp:32:16: warning: 'tree::L' will be initialized after [-Wreorder]
   32 |     vector<vi> L;
      |                ^
split.cpp:31:8: warning:   'vi tree::sz' [-Wreorder]
   31 |     vi sz;
      |        ^~
split.cpp:34:5: warning:   when initialized here [-Wreorder]
   34 |     tree(int N) : n(N), L(N), sz(N, 0) {}
      |     ^~~~
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:2:
split.cpp: In function 'vi find_split(int, int, int, int, vi, vi)':
split.cpp:198:27: warning: comparison of integer expressions of different signedness: 'std::set<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  198 |     assert(NoduriB.size() == Mar[1].first);
      |            ~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~
split.cpp:199:27: warning: comparison of integer expressions of different signedness: 'std::set<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  199 |     assert(NoduriC.size() == Mar[2].first);
      |            ~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~

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