#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>
#include <functional>
// Define a struct to hold coupon information along with its original index.
struct Coupon {
int id;
long long p; // price
int t; // type
};
// Comparison function for sorting multiplier coupons (T > 1).
// Sorts by P*T/(T-1) asc. using 128-bit cross-multiplication to avoid precision errors.
bool compareMultipliers(const Coupon& a, const Coupon& b) {
return static_cast<unsigned __int128>(a.p) * a.t * (b.t - 1) < static_cast<unsigned __int128>(b.p) * b.t * (a.t - 1);
}
// Comparison function for sorting spender coupons (T = 1) by price.
bool compareT1s(const Coupon& a, const Coupon& b) {
return a.p < b.p;
}
std::vector<int> max_coupons(int A_int, std::vector<int> P, std::vector<int> T) {
long long N = P.size();
unsigned __int128 A = A_int;
// 1. Separate coupons into multipliers and spenders.
std::vector<Coupon> multipliers;
std::vector<Coupon> t1s;
for (int i = 0; i < N; ++i) {
if (T[i] == 1) {
t1s.push_back({i, (long long)P[i], T[i]});
} else {
multipliers.push_back({i, (long long)P[i], T[i]});
}
}
// Sort both groups based on the optimal buying strategy.
std::sort(multipliers.begin(), multipliers.end(), compareMultipliers);
std::sort(t1s.begin(), t1s.end(), compareT1s);
// 2. Precompute prefix sums for spender prices for efficient lookups.
std::vector<unsigned __int128> t1_prefix_sum(t1s.size() + 1, 0);
for (size_t i = 0; i < t1s.size(); ++i) {
t1_prefix_sum[i + 1] = t1_prefix_sum[i] + t1s[i].p;
}
// Helper to count affordable spenders using binary search on prefix sums.
auto count_affordable_t1s = [&](unsigned __int128 current_tokens) {
if (current_tokens == 0) return 0;
auto it = std::upper_bound(t1_prefix_sum.begin() + 1, t1_prefix_sum.end(), current_tokens);
return static_cast<int>(std::distance(t1_prefix_sum.begin(), it) - 1);
};
int max_total_coupons = 0;
int best_k_m = 0; // The optimal number of multipliers to buy.
// 3. Baseline: buy zero multipliers.
max_total_coupons = count_affordable_t1s(A);
// A token amount large enough to afford essentially all remaining coupons.
// Sum of all prices is a safe threshold. (N_max * P_max ≈ 2e5 * 1e9 = 2e14)
const unsigned __int128 THRESHOLD = (unsigned __int128)200001 * 1000000000;
unsigned __int128 current_tokens = A;
// 4. Iterate through multiplier prefixes.
for (int k_m = 0; k_m < multipliers.size(); ++k_m) {
const auto& m = multipliers[k_m];
if (current_tokens >= m.p) {
// "Buy" the multiplier and update tokens.
current_tokens = (current_tokens - m.p) * m.t;
int num_t1s;
// Optimization: if tokens are huge, we can afford all spenders.
if (current_tokens >= THRESHOLD) {
num_t1s = t1s.size();
} else {
num_t1s = count_affordable_t1s(current_tokens);
}
// Check if this new combination of (k_m+1) multipliers is better.
if (k_m + 1 + num_t1s > max_total_coupons) {
max_total_coupons = k_m + 1 + num_t1s;
best_k_m = k_m + 1;
}
} else {
// If we can't afford this multiplier, we can't buy any longer prefixes.
break;
}
}
// 5. Reconstruct the final answer.
std::vector<int> result;
result.reserve(max_total_coupons);
// Add the best prefix of multipliers.
for (int i = 0; i < best_k_m; ++i) {
result.push_back(multipliers[i].id);
}
// Recalculate the token amount after buying these multipliers.
unsigned __int128 tokens_for_t1s = A;
for (int i = 0; i < best_k_m; ++i) {
tokens_for_t1s = (tokens_for_t1s - multipliers[i].p) * multipliers[i].t;
}
// Determine how many spenders can be bought.
int num_t1s_to_add = count_affordable_t1s(tokens_for_t1s);
// Add the affordable spenders.
for (int i = 0; i < num_t1s_to_add; ++i) {
result.push_back(t1s[i].id);
}
return result;
}
