#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>
// Function to find the maximum number of coupons that can be bought.
// A: Initial amount of tokens.
// P: Vector of coupon prices.
// T: Vector of coupon multipliers.
// Returns a vector of 0-indexed coupon IDs in the optimal buying order.
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
struct Coupon {
int p, t, id;
};
int n = P.size();
if (n == 0) {
return {};
}
std::vector<Coupon> coupons(n);
for (int i = 0; i < n; ++i) {
coupons[i] = {P[i], T[i], i};
}
// Sort coupons based on the optimal buying order criterion.
// The lambda function implements the comparison logic.
std::sort(coupons.begin(), coupons.end(), [](const Coupon& a, const Coupon& b) {
if (a.t == 1 && b.t == 1) {
if (a.p != b.p) return a.p < b.p;
return a.id < b.id; // Stable tie-breaking
}
if (a.t == 1) return false; // Coupons with T=1 come last.
if (b.t == 1) return true;
// For T > 1, compare Pa*Ta/(Ta-1) and Pb*Tb/(Tb-1).
// To avoid floating point, we cross-multiply:
// Pa*Ta*(Tb-1) vs Pb*Tb*(Ta-1)
long long val_a = 1LL * a.p * a.t * (b.t - 1);
long long val_b = 1LL * b.p * b.t * (a.t - 1);
if (val_a != val_b) {
return val_a < val_b;
}
return a.id < b.id; // Stable tie-breaking
});
// Calculate suffix sums of prices to handle token overflow.
std::vector<long long> suffix_p(n + 1, 0);
for (int i = n - 1; i >= 0; --i) {
suffix_p[i] = suffix_p[i + 1] + coupons[i].p;
}
long long current_tokens = A;
std::vector<int> bought_coupons_indices;
for (int i = 0; i < n; ++i) {
// Optimization: If current tokens are enough to cover all remaining prices,
// we can afford to buy all of them. This avoids token overflow.
if (current_tokens >= suffix_p[i]) {
for (int j = i; j < n; ++j) {
bought_coupons_indices.push_back(coupons[j].id);
}
break;
}
if (current_tokens >= coupons[i].p) {
current_tokens -= coupons[i].p;
current_tokens *= coupons[i].t;
bought_coupons_indices.push_back(coupons[i].id);
} else {
// Cannot afford the current coupon, so cannot afford any subsequent ones
// in the sorted list.
break;
}
}
return bought_coupons_indices;
}
