#include "prize.h"
#include <bits/stdc++.h>
int lm(int x) {
std::vector<int> val;
val = {1};
while (true) {
const int a = val.back();
val.push_back(a * a + 1);
if (std::accumulate(val.begin(), val.end(), 0) > x) {
break;
}
}
val.pop_back();
if (val.size() == 1) {
return 0;
}
while (true) {
val[val.size() - 2]++;
val[val.size() - 1] = val[val.size() - 2] * val[val.size() - 2] + 1;
if (std::accumulate(val.begin(), val.end(), 0) > x) {
val[val.size() - 2]--;
return x - val[val.size() - 2] * val[val.size() - 2] - 1;
}
}
return -1;
}
std::map<int, std::vector<int>> askMemo;
std::vector<int> Ask(int i) {
if (askMemo.find(i) != askMemo.end()) return askMemo[i];
return askMemo[i] = ask(i);
}
std::vector<int> findLower(std::vector<int> ids) {
auto question = [&](const int i) {
return Ask(ids[i]);
};
const int n = (int)ids.size();
const int limit = lm(n);
std::vector<std::vector<int>> vals;
for (int i = 0; i <= limit; ++i) {
vals.push_back(question(i));
}
int df = 0;
for (const auto &vec : vals) {
df = std::max(df, vec[0] + vec[1]);
}
std::vector<int> lws;
auto solve = [&](auto &&self, const int l, const int r, const int lv, const int rv) -> void {
const int mid = (l + r) / 2;
int lm = -1;
for (int i = mid; i < r; ++i) {
auto x = question(i);
if (x[0] + x[1] != df) {
lws.push_back(ids[i]);
} else {
lm = i;
break;
}
}
int rm = n + 1;
if (lm == mid) rm = mid;
if (rm == mid) {
auto x = question(rm);
// relax(rm, x);
const int newLv = lv, newRv = x[1];
if (newLv + newRv != df) {
self(self, l, rm, newLv, newRv);
}
} else if (mid - l > 1) {
int newLv = lv;
int newRv = rv;
if (lm != -1) {
auto x = question(lm);
newRv = x[1] + (lm - mid);
} else {
newRv = rv + (r - mid);
}
if (newLv + newRv != df) {
self(self, l, mid, newLv, newRv);
}
}
if (lm != -1) {
auto x = question(lm);
const int newLv = x[0], newRv = rv;
if (newLv + newRv != df) {
self(self, lm + 1, r, newLv, newRv);
}
}
};
solve(solve, 0, n, 0, 0);
std::sort(lws.begin(), lws.end());
return lws;
}
int find_best(int n) {
std::vector<int> ids(n);
std::iota(ids.begin(), ids.end(), 0);
while (ids.size() != 1) {
ids = findLower(ids);
}
return ids[0];
}
