IOI '20 - Mushrooms (87.5969pts)

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

using namespace std;

int count_mushrooms(int n) {
    int s = 141; // sqrt(20000)
    // ADD RANDOM_SHUFFLE LATER

    // find first 2s values pairwise
    vector<int> id[2]; // A, B
    id[0].push_back(0);
    vector<int> m;
    int last = min(2*s, n);
    for (int i = 1; i < last; ) {
        if (i+1 < last && (id[0].size() >= 2 || id[1].size() >= 2)) {
            int same = 0;
            if (id[0].size() < 2) same = 1;

            m = {id[same][0], i, id[same][1], i+1};
            int x = use_machine(m);
            if (x & 1) id[1-same].push_back(i+1);
            else id[same].push_back(i+1);
            if (x & 2) id[1-same].push_back(i);
            else id[same].push_back(i);
            i += 2;
        } else {
            m = {0, i};
            if (use_machine(m)) id[1].push_back(i);
            else id[0].push_back(i);
            i++;
        }
    }

    // find the rest of the values as blocks of s
    int A = id[0].size();
    for (int i = last; i < n; ) {
        int big = (id[1].size() > id[0].size() ? 1 : 0);
        int cnt = min((int)id[big].size(), n-i);
        m.clear();
        for (int j = 0; j < cnt; ++j) {
            m.push_back(i+j);
            m.push_back(id[big][j]);
        }
        int x = use_machine(m);
        if (x & 1) {
            id[1-big].push_back(i);
            ++x;
        } else id[big].push_back(i);
        x /= 2;
        if (big == 0) x = cnt-x;
        A += x;
        i += cnt;
    }
    return A;
}