#include "souvenirs.h"
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define vll vector<ll>
#define pll pair<ll, ll>
void buy_souvenirs(int n, long long P0) {
map<ll, pair<vector<int>, ll>> mp;
vll costs(n, 0), bought(n, 0);
costs[0] = P0;
vll queried;
for (int i = 0; i < n; i++) {
if (!costs[i]) {
vector<pair<set<int>, ll>> eqs;
pair<vector<int>, ll> ret;
if (mp.count(costs[i-1]-1)) ret = mp[costs[i-1]-1];
else {
ret = transaction(costs[i-1]-1);
for (int x : ret.first) bought[x]++;
mp[costs[i-1]-1] = ret;
}
set<int> ss;
for (int x : ret.first) ss.insert(x);
eqs.push_back({ss, costs[i-1]-1-ret.second});
//deal with the last equation we have left, query its average
while (eqs.size()) {
//simplify equation just in case
set<int> torem;
for (int x : eqs.back().first) if (costs[x]) {
torem.insert(x);
eqs.back().second -= costs[x];
}
for (ll x : torem) eqs.back().first.erase(x);
vector<int> pseud;
for (int xx : eqs.back().first) pseud.push_back(xx);
mp[eqs.back().second] = make_pair(pseud, 0LL);
//if 1 element we can directly find out what it is
if (eqs.back().first.size()==1) {
ll x = *eqs.back().first.begin();
costs[x] = eqs.back().second;
eqs.pop_back();
//check if we know anything smaller
bool cont = false;
for (int j = x+1; j < n; j++) {
if (!costs[j]) {
cont = true;
x = j-1;
break;
}
}
if (cont) {
pair<vector<int>, ll> pr;
if (mp.count(costs[x]-1)) pr = mp[costs[x]-1];
else {
pr = transaction(costs[x]-1);
for (int x : pr.first) bought[x]++;
mp[costs[x]-1] = pr;
}
set<int> s;
for (int y : pr.first) s.insert(y);
eqs.push_back({s, costs[x] - 1 - pr.second});
}
} else if (!eqs.back().first.size()) eqs.pop_back();
else {
//otherwise we want to query the average of its cost
pair<vector<int>, ll> pr;
if (mp.count(eqs.back().second / eqs.back().first.size())) {
pr = mp[eqs.back().second / eqs.back().first.size()];
} else {
pr = transaction(eqs.back().second / eqs.back().first.size());
for (int x : pr.first) bought[x]++;
mp[eqs.back().second / eqs.back().first.size()] = pr;
}
set<int> s;
for (int x : pr.first) s.insert(x);
eqs.push_back({s, eqs.back().second / eqs.back().first.size() - pr.second});
}
}}
}
for (int i = 0; i < n; i++) {
while (bought[i] < i) {
transaction(costs[i]);
bought[i]++;
}
}
return;
}
