#include "mushrooms.h"
#include <bits/stdc++.h>
using namespace std;
#define ll long long int
#define vll vector<ll>
#define vi vector<int>
#define pi pair<int, int>
#define pll pair<ll, ll>
#define all(x) (x).begin(), (x).end()
#define fori(i,n) for(int i = 0; i < int(n); ++i)
bool test_is_a(int i) {
return use_machine(vi{0, i}) == 0;
}
int count_mushrooms(int n) {
vi opts(n-1);
fori(i, n-1) opts[i] = i + 1;
random_shuffle(all(opts));
auto pop1 = [&]() {
int i = opts.back();
opts.pop_back();
return i;
};
auto pop2 = [&]() -> pi { return {pop1(), pop1()}; };
vi as = {0};
vi bs = {};
auto add_to_list = [&](int idx, bool is_a) {
if(is_a) as.push_back(idx);
else bs.push_back(idx);
};
auto biggest_list_size = [&]() { return max(as.size(), bs.size()); };
auto get_bigger_list = [&]() -> pair<vi&, bool> {
if(as.size() >= bs.size()) return {as, true};
else return {bs, false};
};
if(n <= 226) {
fori(_i, n-1) {
int i = pop1();
add_to_list(i, test_is_a(i));
}
return as.size();
}
while (biggest_list_size() < 2) {
int i = pop1();
add_to_list(i, test_is_a(i));
}
while (biggest_list_size() < 76) {
auto [i1, i2] = pop2();
const auto& [ls, use_a] = get_bigger_list();
vi test = {i1, ls[0], i2, ls[1]};
int res = use_machine(test);
add_to_list(i1, ((res&1) == 1) ^ use_a);
add_to_list(i2, ((res&2) == 2) ^ use_a);
}
int as_count = 0;
int bs_count = 0;
while (!opts.empty()) {
const auto& [ls, use_a] = get_bigger_list();
vi test;
for(int i : ls) {
if(opts.empty()) break;
test.push_back(pop1());
test.push_back(i);
}
int res = use_machine(test);
add_to_list(test[0], (res&1) ^ use_a);
int bc1 = res >> 1;
int ac1 = test.size()/2 - 1 - bc1;
if (!use_a) swap(ac1, bc1);
as_count += ac1;
bs_count += bc1;
}
as_count += as.size();
bs_count += bs.size();
return as_count;
}
