#include <vector>
#include <numeric>
#include <cmath>
#include <algorithm>
#include <iostream>
#include "souvenirs.h"
// Defined by the grader
void buy_souvenirs(int N, long long P0) {
std::vector<long long> P(N, 0);
std::vector<int> current_counts(N, 0);
P[0] = P0;
auto solve_recursive = [&](auto&& self, std::vector<int> U, long long current_sum) -> void {
// Base Case: If only one item, its price is the total sum.
if (U.size() == 1) {
P[U[0]] = current_sum;
return;
}
long long count = U.size();
long long safe_M = current_sum / count;
if (safe_M < 1) safe_M = 1;
// Perform transaction
auto res = transaction(safe_M);
std::vector<int> L_new = res.first;
long long R_new = res.second;
// Track what we bought to ensure we don't violate quotas later
for (int idx : L_new) current_counts[idx]++;
std::vector<int> U_hit;
std::vector<int> U_miss;
// We use a simple boolean map for O(1) lookup of L_new elements
std::vector<bool> in_L_new(N, false);
for (int idx : L_new) in_L_new[idx] = true;
long long sum_hit = safe_M - R_new;
for (int idx : U) {
if (in_L_new[idx]) {
U_hit.push_back(idx);
} else {
U_miss.push_back(idx);
}
}
if (U_hit.empty()) {
return;
}
// Recurse
long long sum_miss = current_sum - sum_hit;
if (!U_hit.empty()) self(self, U_hit, sum_hit);
if (!U_miss.empty()) self(self, U_miss, sum_miss);
};
// Initial Transaction
// Ensure we buy at least one item. P[0]-1 guarantees skipping Type 0.
long long initial_M = P0 - 1;
auto initial_res = transaction(initial_M);
std::vector<int> initial_L = initial_res.first;
long long initial_R = initial_res.second;
for (int idx : initial_L) current_counts[idx]++;
long long total_sum = initial_M - initial_R;
solve_recursive(solve_recursive, initial_L, total_sum);
for (int i = 1; i < N; ++i) {
int target = i;
while (current_counts[i] < target) {
auto res = transaction(P[i]);
for (int bought : res.first) {
current_counts[bought]++;
}
}
}
}
