#include "souvenirs.h"
#include <vector>
#include <algorithm>
using namespace std;
void buy_souvenirs(int N, long long P0) {
int n_types = N - 1; // types 1 .. N-1
vector<long long> price(N, 0); // price[i], 1-indexed
vector<int> cnt(N, 0); // how many of each type already bought
vector<bool> known(N, false); // whether price[i] is already known
struct Equation {
vector<int> indices;
long long sum;
};
vector<Equation> equations; // all equations obtained so far
// after adding a new equation, try to deduce all possible exact prices
auto deduce = [&]() {
bool changed;
do {
changed = false;
for (auto& eq : equations) {
// count unknown indices in this equation
long long known_sum = 0;
int unknown_idx = -1;
int unknown_cnt = 0;
for (int i : eq.indices) {
if (known[i]) {
known_sum += price[i];
} else {
++unknown_cnt;
unknown_idx = i;
}
}
if (unknown_cnt == 1) {
long long candidate = eq.sum - known_sum;
if (candidate >= 1 && !known[unknown_idx]) {
price[unknown_idx] = candidate;
known[unknown_idx] = true;
changed = true;
}
}
}
} while (changed);
};
// perform discovery queries until all prices are known
int remaining = n_types;
while (remaining > 0) {
// find the largest index still unknown (the smallest price)
int k = -1;
for (int i = n_types; i >= 1; --i)
if (!known[i]) { k = i; break; }
// compute a tight upper bound for price[k]
long long UB = P0 - k; // global bound
// use already known prices (they are larger than k -> give an upper bound)
for (int i = 1; i < k; ++i) {
if (known[i]) {
// price[k] <= price[i] - (k - i)
long long cand = price[i] - (k - i);
if (cand < UB) UB = cand;
}
}
// use all equations
for (const auto& eq : equations) {
long long adj_sum = eq.sum;
vector<int> unk;
for (int i : eq.indices) {
if (known[i]) {
adj_sum -= price[i];
} else {
unk.push_back(i);
}
}
if (unk.empty()) continue;
// all indices in unk are <= k (k is the largest unknown)
long long offset = 0;
for (int i : unk) offset += (k - i);
int cnt_unk = unk.size();
if (offset > adj_sum) continue;
long long cand = (adj_sum - offset) / cnt_unk;
if (cand < 1) cand = 1;
if (cand < UB) UB = cand;
}
if (UB < 1) UB = 1;
long long M = UB; // safe and optimal next query
auto res = transaction(M);
const vector<int>& L = res.first;
long long R = res.second;
// update counts
for (int i : L) cnt[i]++;
// add the equation from this query and propagate deductions
Equation eq;
eq.indices = L;
eq.sum = M - R;
equations.push_back(eq);
deduce();
// count still unknown types
remaining = 0;
for (int i = 1; i <= n_types; ++i) if (!known[i]) ++remaining;
}
// now all prices are known, buy the remaining copies one by one
for (int i = 1; i <= n_types; ++i) {
int need = i - cnt[i];
while (need > 0) {
transaction(price[i]); // buys exactly one souvenir of type i
++cnt[i];
--need;
}
}
}
