#include "souvenirs.h"
#include <bits/stdc++.h>
using namespace std;
using int64 = long long;
void buy_souvenirs(int N, long long P0) {
vector<long long> price(N, -1);
price[0] = P0;
// depend[i]: one equation for unknown prices in this set
// sumv[i] = sum of prices in depend[i]
vector<set<int>> depend(N);
vector<long long> sumv(N, 0);
// count how many souvenirs of each type have already been bought
// during the discovery phase
vector<int> bought(N, 0);
auto apply_query = [&](long long M, set<int>& unknown_set, long long& unknown_sum) {
auto res = transaction(M);
const vector<int>& L = res.first;
long long R = res.second;
set<int> s;
for (int x : L) {
s.insert(x);
bought[x]++;
}
long long spent = M - R;
// remove already known prices from the equation
for (int i = 0; i < N; ++i) {
if (price[i] != -1 && s.count(i)) {
spent -= price[i];
s.erase(i);
}
}
unknown_set = s;
unknown_sum = spent;
};
while (true) {
int left = 0;
for (int i = 0; i < N; ++i) {
if (price[i] == -1) left++;
}
if (left == 0) break;
bool progressed = false;
bool have_unknown_suffix = false;
for (int i = N - 1; i >= 0; --i) {
if (price[i] == -1) have_unknown_suffix = true;
// From a known price price[i], query price[i]-1 to get
// an equation involving only cheaper unknown items.
if (price[i] != -1 && have_unknown_suffix) {
set<int> s;
long long sm = 0;
apply_query(price[i] - 1, s, sm);
depend[i + 1] = s;
sumv[i + 1] = sm;
progressed = true;
break;
}
// If an equation contains exactly one unknown, we recover it.
if (!depend[i].empty() && depend[i].size() == 1) {
int x = *depend[i].begin();
price[x] = sumv[i];
// substitute into all stored equations
for (int j = 0; j < N; ++j) {
if (depend[j].count(x)) {
depend[j].erase(x);
sumv[j] -= price[x];
}
}
depend[i].clear();
sumv[i] = 0;
progressed = true;
break;
}
// Otherwise, split the equation by querying the average.
if (!depend[i].empty()) {
long long x = sumv[i] / (long long)depend[i].size();
set<int> s;
long long sm = 0;
apply_query(x, s, sm);
int k = *s.begin();
depend[k] = s;
sumv[k] = sm;
progressed = true;
break;
}
}
if (!progressed) break;
}
// Top up counts to exactly i souvenirs of type i.
for (int i = 1; i < N; ++i) {
while (bought[i] < i) {
transaction(price[i]);
bought[i]++;
}
}
}
