#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>
#include <climits> // For LLONG_MAX
// Use a struct to keep coupon properties and original index together.
struct Coupon {
int id;
long long p; // Use long long for price to handle large numbers.
int t;
};
/**
* @brief Solves the coupon buying problem to maximize the number of coupons purchased.
*
* @param A_int The initial number of tokens.
* @param P A vector of coupon prices.
* @param T A vector of coupon multipliers.
* @return A vector of 0-indexed coupon IDs in the optimal buying order.
*/
std::vector<int> max_coupons(int A_int, std::vector<int> P, std::vector<int> T) {
long long A = A_int;
int N = P.size();
// 1. Create a vector of Coupon structs to store all properties.
std::vector<Coupon> coupons(N);
for (int i = 0; i < N; ++i) {
coupons[i] = {i, (long long)P[i], T[i]};
}
// 2. Sort coupons using the custom greedy comparator.
std::sort(coupons.begin(), coupons.end(), [](const Coupon& c1, const Coupon& c2) {
// T=1 coupons are always worse; they are processed after T>1 coupons.
if (c1.t == 1 && c2.t > 1) return false; // c2 comes first
if (c1.t > 1 && c2.t == 1) return true; // c1 comes first
// If both are T=1, cheaper ones are better.
if (c1.t == 1 && c2.t == 1) {
return c1.p < c2.p;
}
// If both have T > 1, compare their "value" P*T/(T-1).
// Using long double is a standard-compliant way to get high precision.
long double v1 = static_cast<long double>(c1.p) * c1.t / (c1.t - 1);
long double v2 = static_cast<long double>(c2.p) * c2.t / (c2.t - 1);
// Tie-break with price to ensure a consistent, stable order.
if (v1 != v2) {
return v1 < v2;
}
return c1.p < c2.p;
});
// 3. Binary search for the maximum number of coupons (k).
auto check = [&](int k) {
if (k == 0) return true;
long long current_tokens = A;
for (int i = 0; i < k; ++i) {
long long price = coupons[i].p;
int mult = coupons[i].t;
if (current_tokens < price) {
return false;
}
current_tokens -= price;
if (mult > 1) {
// Check for impending overflow before multiplication.
if (current_tokens > LLONG_MAX / mult) {
// If it would overflow, tokens become effectively infinite.
// LLONG_MAX is a safe sentinel value for this purpose.
current_tokens = LLONG_MAX;
} else {
current_tokens *= mult;
}
}
}
return true;
};
int low = 0, high = N;
int best_k = 0;
while (low <= high) {
int mid = low + (high - low) / 2;
if (check(mid)) {
best_k = mid;
low = mid + 1;
} else {
high = mid - 1;
}
}
// 4. Construct the result vector with the IDs of the first `best_k` coupons.
std::vector<int> result;
result.reserve(best_k);
for (int i = 0; i < best_k; ++i) {
result.push_back(coupons[i].id);
}
return result;
}
