#include "souvenirs.h"
#include <vector>
#include <numeric>
#include <algorithm>
#include <map>
void buy_souvenirs(int N, long long P0) {
std::vector<long long> p(N);
p[0] = P0;
std::map<int, int> counts;
long long current_sum = 0;
// Phase 1: Price Discovery
for (int i = 1; i < N; ++i) {
current_sum += p[i - 1];
long long low = 1, high = p[i - 1] - 1;
long long found_p = p[i - 1];
while (low <= high) {
long long mid = low + (high - low) / 2;
if (mid == 0) {
low = 1;
continue;
}
auto result = transaction(current_sum + mid);
for(int item : result.first) {
counts[item]++;
}
bool bought_i = false;
// The returned vector is sorted, so we can check efficiently.
for (int item : result.first) {
if (item == i) {
bought_i = true;
break;
}
if (item > i) {
break;
}
}
if (bought_i) {
found_p = mid;
high = mid - 1;
} else {
low = mid + 1;
}
}
p[i] = found_p;
}
// Phase 2: Purchasing missing souvenirs
// We iterate from N-1 down to 0 to ensure that when we buy item i,
// we don't accidentally buy any other items we still need.
for (int i = N - 1; i >= 0; --i) {
if (counts.find(i) == counts.end() || counts[i] < 1) {
// A transaction with p[i] coins will buy souvenir i and nothing else.
auto res = transaction(p[i]);
for(int item : res.first) {
counts[item]++;
}
}
}
}
