#include "souvenirs.h"
#include <map>
#include <set>
#include <utility>
#include <vector>
using namespace std;
namespace {
struct Equation {
bool has = false;
vector<int> coeff;
long long rhs = 0;
};
} // namespace
void buy_souvenirs(int N, long long P0) {
vector<int> cnt(N, 0);
vector<int> target(N, 0);
for (int i = 0; i < N; ++i) target[i] = i;
long long calls = 0;
bool overshot = false;
map<long long, vector<int>> known_types;
vector<Equation> equations(N);
vector<long long> price(N, -1);
price[0] = P0;
auto done = [&]() -> bool {
for (int i = 1; i < N; ++i) {
if (cnt[i] != target[i]) return false;
}
return true;
};
auto safe_known_type = [&](const vector<int> &L) -> bool {
for (int idx : L) {
if (cnt[idx] >= target[idx]) return false;
}
return true;
};
auto do_transaction = [&](long long M) -> pair<vector<int>, long long> {
pair<vector<int>, long long> res = transaction(M);
++calls;
const vector<int> &L = res.first;
for (int idx : L) {
++cnt[idx];
if (cnt[idx] > target[idx]) overshot = true;
}
if (known_types.find(M) == known_types.end()) {
known_types[M] = L;
}
if (!L.empty()) {
int first = L.front();
if (first >= 0 && first < N && !equations[first].has) {
equations[first].has = true;
equations[first].rhs = M - res.second;
equations[first].coeff.assign(N, 0);
for (int idx : L) equations[first].coeff[idx] = 1;
}
}
return res;
};
auto solve_price_closure = [&]() -> bool {
bool changed = false;
for (int i = N - 1; i >= 1; --i) {
if (price[i] != -1) continue;
if (!equations[i].has) continue;
long long v = equations[i].rhs;
bool ready = true;
for (int j = i + 1; j < N; ++j) {
if (!equations[i].coeff[j]) continue;
if (price[j] == -1) {
ready = false;
break;
}
v -= price[j];
}
if (!ready) continue;
if (v <= 0) continue;
if (price[i - 1] != -1 && v >= price[i - 1]) continue;
if (i + 1 < N && price[i + 1] != -1 && v <= price[i + 1]) continue;
price[i] = v;
changed = true;
}
return changed;
};
if (N == 2) {
if (calls < 5000) do_transaction(P0 - 1);
return;
}
// Phase 1: one descending discovery chain.
long long M = P0 - 1;
set<long long> seen;
const int max_probe = min(400, 4 * N);
for (int step = 0; step < max_probe; ++step) {
if (calls >= 5000 || overshot || done()) break;
if (M <= 0 || M >= P0) break;
if (seen.find(M) != seen.end()) break;
seen.insert(M);
auto res = do_transaction(M);
const vector<int> &L = res.first;
long long R = res.second;
if (L.empty()) break;
if (L.front() == N - 1 && (int)L.size() == 1) break;
long long spent = M - R;
long long nextM;
if ((int)L.size() == 1) {
nextM = spent - 1;
} else {
long long avg = spent / (long long)L.size();
if (L.front() == 1) {
nextM = avg;
} else {
nextM = (M + avg) / 2;
}
}
if (nextM >= M) nextM = M - 1;
if (nextM < 1) nextM = 1;
M = nextM;
}
while (solve_price_closure()) {
}
// Phase 2: if P[i-1] is known and row i is missing, query M=P[i-1]-1.
bool phase2_progress = true;
while (phase2_progress && calls < 5000 && !overshot && !done()) {
while (solve_price_closure()) {
}
phase2_progress = false;
int need = -1;
for (int i = 2; i < N; ++i) {
if (!equations[i].has && price[i - 1] != -1) {
need = i;
break;
}
}
if (need != -1) {
long long q = price[need - 1] - 1;
if (q >= 1 && q < P0) {
do_transaction(q);
phase2_progress = true;
}
}
}
while (solve_price_closure()) {
}
bool all_prices_known = true;
for (int i = 1; i < N; ++i) {
if (price[i] == -1) {
all_prices_known = false;
break;
}
}
// Exact singleton completion if prices are known.
if (all_prices_known && !overshot) {
for (int i = 1; i < N && calls < 5000 && !overshot; ++i) {
while (cnt[i] < target[i] && calls < 5000 && !overshot) {
do_transaction(price[i]);
}
}
return;
}
// Fallback: use only previously observed safe transaction types.
while (calls < 5000 && !overshot && !done()) {
long long bestM = -1;
long long bestScore = -1;
for (const auto &it : known_types) {
const long long candM = it.first;
const vector<int> &L = it.second;
if (!safe_known_type(L)) continue;
long long score = 0;
for (int idx : L) score += (long long)(target[idx] - cnt[idx]);
if (score > bestScore) {
bestScore = score;
bestM = candM;
}
}
if (bestM == -1) break;
do_transaction(bestM);
}
}
