#include "mushrooms.h"
#include <algorithm>
#include <random>
using namespace std;
/*
2 * START + (n - 2 * START) / (START - 1)
*/
int count_mushrooms(int n) {
vector<int> knownAs, knownBs;
knownAs.push_back(0);
int pos = 1;
int res = knownAs.size();
vector<int> perm(n);
for (int i = 0; i < n; i++) perm[i] = i;
mt19937 rng(7);
shuffle(perm.begin() + 1, perm.end(), rng);
while (pos < n) {
int querySize = max(knownAs.size(), knownBs.size());
int mushroomsInQuery = min(querySize, n - pos);
vector<int> query;
bool isA = knownAs.size() >= knownBs.size();
for (int i = 0; i < mushroomsInQuery; i++) {
if (isA) {
query.push_back(knownAs[i]);
} else {
query.push_back(knownBs[i]);
}
query.push_back(perm[pos++]);
}
int qres = use_machine(query);
int numChanges = qres / 2 + (qres % 2);
if (isA) {
res += mushroomsInQuery - numChanges;
if (qres % 2 == 0) knownAs.push_back(perm[pos-1]);
else knownBs.push_back(perm[pos-1]);
} else {
res += numChanges;
if (qres % 2 == 0) knownBs.push_back(perm[pos-1]);
else knownAs.push_back(perm[pos-1]);
}
}
return res;
}
/*
B_B_B_B_B_
*/
