#include "souvenirs.h"
#include <utility>
#include <vector>
#include <iostream>
using namespace std;
using ll = long long;
void buy_souvenirs(int N, ll P0) {
// history[first] = { indices, cost }
// if cost == 0, then unknown / irrelevant
vector<pair<vector<int>, ll>> history(N);
vector<ll> value(N);
vector<int> count(N);
auto do_trans = [&](ll paid) -> int {
auto [indices, change] = transaction(paid);
for (int i : indices) count[i]++;
history[indices[0]] = { indices, paid - change };
return indices[0];
};
auto eval = [&](vector<int> indices, ll cost, int cur) -> pair<int, ll> {
int unknown_cnt = 0;
for (int i : indices) {
if (i > cur) cost -= value[i];
else unknown_cnt++;
}
return { unknown_cnt, cost };
};
history[0] = { { 0 }, P0 };
for (int cur = N-1; cur > 0; cur--) {
// find last known
int last_known = cur;
while (history[last_known].second == 0) last_known--;
// go down until cur
while (last_known < cur) {
auto [indices, total_cost] = history[last_known];
auto [unknown_cnt, cost] = eval(indices, total_cost, cur);
if (unknown_cnt == 1) last_known = do_trans(cost - 1);
else last_known = do_trans(cost / unknown_cnt);
}
// history[cur] = { { cur }, cur_cost }
auto [indices, total_cost] = history[last_known];
value[cur] = eval(indices, total_cost, cur).second;
history[cur].second = 0;
}
// fill up the staircase
for (int i = 1; i < N; i++) {
while (count[i] != i) do_trans(value[i]);
}
}
