Submission #1077820

#	Time	Username	Problem	Language	Result	Execution time	Memory
1077820		mc061	Split the Attractions (IOI19_split)	C++17	40 / 100	843 ms	18716 KiB

split

#include <bits/stdc++.h>
using namespace std;

const int N = 1e5+5;
vector<int> tree[N];
int subtreesize[N]={};
int par[N]={};
vector<int> ans;

const int max_afford = 1e7;

int go(int v, int p=0) {
    par[v]=p;
    for (int i = 0; i < tree[v].size(); ++i) {
        int u = tree[v][i];
        if (u != p) subtreesize[v] += go(u, v);
    }
    return subtreesize[v];
}
int cursz;
void fill_ans(int v, int p, int num) {
    if (ans[v] != 0) return;
    if (cursz == 0) return;
    cursz--;
    ans[v] = num;
    for (int i = 0; i < tree[v].size(); ++i) {
        int u = tree[v][i];
        if (u != p) fill_ans(u, v, num);
    }
}

bool solve_for_tree(vector<pair<int, int>> edges, int sz1, int sz2, int num1, int num2) {
    const int n = edges.size()+1;
    for (int i = 0; i < n; ++i) tree[i].clear(), subtreesize[i]=1;
    for (int i = 0; i < n-1; ++i) {
        tree[edges[i].first].push_back(edges[i].second);
        tree[edges[i].second].push_back(edges[i].first);
    }
    go(0);
    for (int i = 0; i < n; ++i) {
        int here = subtreesize[i];
        int rest = n-subtreesize[i];
        if (here >= sz1 && rest >= sz2) {
            cursz = sz1;
            fill_ans(i, par[i], num1);
            cursz = sz2;
            fill_ans(0, 0, num2);
            for (int j = 0; j < n; ++j) {
                if (ans[j] == 0) ans[j] = 6 - num1 - num2;
            }
            return true;
        }
        if (here >= sz2 && rest >= sz1) {
            cursz = sz2;
            fill_ans(i, par[i], num2);
            cursz = sz1;
            fill_ans(0, 0, num1);
            for (int j = 0; j < n; ++j) {
                if (ans[j] == 0) ans[j] = 6 - num1 - num2;
            }
            return true;
        }
    }
    return false;
}

vector<pair<int, int>> e;
vector<pair<int, int>> all_e;

struct DSU {
    vector<int> p;
    vector<int> sz;
    void init(int n) {
        p.resize(n);
        sz.resize(n);
        p.assign(n, 1);
        sz.assign(n, 1);
        iota(p.begin(), p.end(), 0);
    }
    DSU(int n) {
        init(n);
    }
    int get(int a) {
        return p[a] = (a == p[a] ? a : get(p[a]));
    }
    void merge(int a, int b) {
        int x = get(a), y = get(b);
        if (x == y) return;
        if (sz[x] > sz[y]) swap(x, y);
        p[x] = y;
        sz[y] += sz[x];
    }
};

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    ans.resize(n, 0);
    int its = max_afford / n;
    for (int i = 0; i < p.size(); ++i) {
        all_e.push_back({p[i], q[i]});
    }
    vector<pair<int, int>> sizes={{a,1}, {b,2}, {c, 3}};
    sort(sizes.begin(), sizes.end());
    DSU d(0);
    for (int _ = 0; _ < its; ++_) {
        d.init(n);
        e.clear();
        e.shrink_to_fit();
        random_shuffle(all_e.begin(), all_e.end());
        int cnt = 0;
        for (int i = 0; i < all_e.size(); ++i) {
            if (d.get(all_e[i].first) != d.get(all_e[i].second)) {
                e.push_back(all_e[i]);
                d.merge(all_e[i].first, all_e[i].second);
                cnt++;
                if (cnt == n-1) break;
            }
        }
        if (solve_for_tree(e, a, b, 1, 2)) break;
    }
    return ans;
}

Compilation message (stderr)

split.cpp: In function 'int go(int, int)':
split.cpp:14:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   14 |     for (int i = 0; i < tree[v].size(); ++i) {
      |                     ~~^~~~~~~~~~~~~~~~
split.cpp: In function 'void fill_ans(int, int, int)':
split.cpp:26:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   26 |     for (int i = 0; i < tree[v].size(); ++i) {
      |                     ~~^~~~~~~~~~~~~~~~
split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:98:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   98 |     for (int i = 0; i < p.size(); ++i) {
      |                     ~~^~~~~~~~~~
split.cpp:110:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  110 |         for (int i = 0; i < all_e.size(); ++i) {
      |                         ~~^~~~~~~~~~~~~~

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