#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>
#include <optional>
// Use a 128-bit integer type for token calculations to prevent overflow.
using int128 = unsigned __int128;
// Coupon struct holding its original index and properties.
struct Coupon {
int id;
long long p;
int t;
};
// Global sorting function for all coupons.
bool compareCoupons(const Coupon& a, const Coupon& b) {
bool a_is_mult = a.t > 1;
bool b_is_mult = b.t > 1;
if (a_is_mult && !b_is_mult) return true; // Multipliers first
if (!a_is_mult && b_is_mult) return false;
if (a_is_mult) { // Both are multipliers
return static_cast<int128>(a.p) * a.t * (b.t - 1) < static_cast<int128>(b.p) * b.t * (a.t - 1);
} else { // Both are spenders
return a.p < b.p;
}
}
std::vector<int> max_coupons(int A_int, std::vector<int> P, std::vector<int> T) {
int N = P.size();
int128 A = A_int;
std::vector<Coupon> sorted_coupons;
for (int i = 0; i < N; ++i) {
sorted_coupons.push_back({i, (long long)P[i], T[i]});
}
std::sort(sorted_coupons.begin(), sorted_coupons.end(), compareCoupons);
// dp[k] = max tokens after buying k coupons. Using a 1D DP array optimized in place.
std::vector<std::optional<int128>> dp(N + 1);
// parent[k] stores the coupon `id` and the previous `k` to reconstruct the path.
std::vector<std::pair<int, int>> parent(N + 1, {-1, -1});
dp[0] = A;
for (const auto& coupon : sorted_coupons) {
for (int k = N - 1; k >= 0; --k) {
if (!dp[k].has_value()) continue;
int128 current_tokens = dp[k].value();
if (current_tokens >= coupon.p) {
int128 new_tokens = (current_tokens - coupon.p) * coupon.t;
if (!dp[k + 1].has_value() || new_tokens > dp[k + 1].value()) {
dp[k + 1] = new_tokens;
// To reconstruct, we need to know we came from state k
// and used this coupon. This is tricky with 1D DP.
// A full 2D parent table is better for clarity.
}
}
}
}
// The 1D DP above finds the max count, but path reconstruction is hard.
// The 2D version below is clearer and allows reconstruction. It will be too slow for large N.
std::vector<std::vector<std::optional<int128>>> dp2D(N + 1, std::vector<std::optional<int128>>(N + 1));
std::vector<std::vector<int>> path_choice(N + 1, std::vector<int>(N + 1, 0)); // 0 = skip, 1 = take
dp2D[0][0] = A;
for (int i = 1; i <= N; ++i) {
const auto& coupon = sorted_coupons[i - 1];
for (int k = 0; k <= i; ++k) {
// Option 1: Don't take the current coupon
dp2D[i][k] = dp2D[i - 1][k];
path_choice[i][k] = 0;
// Option 2: Take the current coupon
if (k > 0) {
if (dp2D[i - 1][k - 1].has_value() && dp2D[i - 1][k - 1].value() >= coupon.p) {
int128 new_tokens = (dp2D[i - 1][k - 1].value() - coupon.p) * coupon.t;
if (!dp2D[i][k].has_value() || new_tokens > dp2D[i][k].value()) {
dp2D[i][k] = new_tokens;
path_choice[i][k] = 1;
}
}
}
}
}
// Find the maximum number of coupons bought
int max_k = 0;
for (int k = N; k >= 0; --k) {
if (dp2D[N][k].has_value()) {
max_k = k;
break;
}
}
// Reconstruct the path by backtracking
std::vector<int> result;
int current_k = max_k;
for (int i = N; i > 0 && current_k > 0; --i) {
if (path_choice[i][current_k] == 1) {
result.push_back(sorted_coupons[i - 1].id);
current_k--;
}
}
// The reconstructed path is in reverse order of processing.
// The problem requires chronological order of purchase.
// Since our processing order *is* the optimal purchase order, we just need to reverse the result.
std::reverse(result.begin(), result.end());
return result;
}
