#include <vector>
#include <algorithm>
#include <iostream>
using namespace std;
// Struct to store coupon details
struct Coupon {
int id;
long long p;
long long t;
};
// Comparator for Multipliers (T > 1)
// Sort based on minimizing the "loss" ratio: P*T / (T-1).
// Logic: If we have coupon A and B, we prefer A before B if:
// P_a * T_a * (T_b - 1) < P_b * T_b * (T_a - 1)
bool compareMultipliers(const Coupon& a, const Coupon& b) {
unsigned __int128 val1 = (unsigned __int128)a.p * a.t * (b.t - 1);
unsigned __int128 val2 = (unsigned __int128)b.p * b.t * (a.t - 1);
if (val1 != val2) return val1 < val2;
return a.p < b.p; // Tie-breaker
}
// Comparator for Spenders (T = 1)
bool compareSpenders(const Coupon& a, const Coupon& b) {
return a.p < b.p;
}
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int n = P.size();
vector<Coupon> multipliers;
vector<Coupon> spenders;
long long total_spender_cost = 0;
long long total_multiplier_cost = 0;
for (int i = 0; i < n; ++i) {
if (T[i] == 1) {
spenders.push_back({i, (long long)P[i], (long long)T[i]});
total_spender_cost += P[i];
} else {
multipliers.push_back({i, (long long)P[i], (long long)T[i]});
total_multiplier_cost += P[i];
}
}
sort(multipliers.begin(), multipliers.end(), compareMultipliers);
sort(spenders.begin(), spenders.end(), compareSpenders);
// Precompute prefix sums of costs for spenders (for binary search)
int s_size = spenders.size();
vector<long long> spender_prefix(s_size + 1, 0);
for (int i = 0; i < s_size; ++i) {
spender_prefix[i+1] = spender_prefix[i] + spenders[i].p;
}
// DP Initialization
int m = multipliers.size();
int limit_k = min(m, 2000); // Check up to 65 multipliers
// dp[k] = max tokens after buying exactly k multipliers from the processed set
vector<long long> dp(limit_k + 1, -1);
// taken[i][k] = true if multiplier i-1 was used to reach state k
// We use a flattened vector or a vector of vectors.
// m can be 100,000, limit_k is 65. Size ~ 6.5MB. Safe.
vector<vector<bool>> taken(m + 1, vector<bool>(limit_k + 1, false));
dp[0] = A;
// We need to track the "Global Best" solution found so far.
// This could be a standard solution (subset of multipliers + subset of spenders)
// OR an "Infinite" solution (subset of multipliers + ALL remaining multipliers + ALL spenders).
struct BestState {
int count;
int k_at_peak; // How many multipliers we picked in DP
int index_at_peak; // Which multiplier index (0..m) we were at
bool is_infinite; // Did we trigger the "Buy All Remaining" condition?
} best_sol = {-1, -1, -1, false};
// Helper to calculate coupons using money
auto check_standard_solution = [&](long long money, int k, int current_idx) {
auto it = upper_bound(spender_prefix.begin(), spender_prefix.end(), money);
int cnt = (int)(it - spender_prefix.begin()) - 1;
int total = k + cnt;
if (total > best_sol.count) {
best_sol = {total, k, current_idx, false};
}
};
// Current cost of multipliers processed so far (prefix sum logic)
long long current_mul_prefix_cost = 0;
// Check initial state (0 multipliers)
check_standard_solution(A, 0, 0);
// Limit value to cap money to prevent overflow (enough to buy everything)
long long INF_CAP = total_spender_cost + total_multiplier_cost + 2000;
for (int i = 0; i < m; ++i) {
long long cost = multipliers[i].p;
long long type = multipliers[i].t;
// Before processing item i, check "Infinite Condition" for existing states
// Cost to buy ALL remaining multipliers (including i) + ALL spenders
long long remaining_mul_cost = total_multiplier_cost - current_mul_prefix_cost;
long long cost_for_everything = remaining_mul_cost + total_spender_cost;
for (int k = 0; k <= limit_k; ++k) {
if (dp[k] == -1) continue;
// Check Infinite Mode
// Note: For multipliers, we pay P but gain money.
// The strict condition to "buy everything" is simply having enough money
// to pay the entry cost of the sequence.
// Since T >= 1, money never decreases below (Start - Cost).
// A safe upper bound is checking if we can afford the sum of prices.
if (dp[k] >= cost_for_everything) {
int total_remaining_mul = m - i; // Indices i to m-1
int total = k + total_remaining_mul + s_size;
if (total > best_sol.count) {
best_sol = {total, k, i, true};
}
}
}
// Update DP table
// Iterate backwards
for (int k = limit_k - 1; k >= 0; --k) {
if (dp[k] != -1) {
if (dp[k] >= cost) {
long long next_tokens = (dp[k] - cost) * type;
if (next_tokens > INF_CAP) next_tokens = INF_CAP;
if (next_tokens > dp[k+1]) {
dp[k+1] = next_tokens;
taken[i+1][k+1] = true;
}
}
}
}
// Update prefix cost
current_mul_prefix_cost += cost;
// Check Standard Solution after this step
for (int k = 1; k <= limit_k; ++k) {
if (dp[k] != -1) {
check_standard_solution(dp[k], k, i + 1);
}
}
}
// Reconstruction
vector<int> result;
// 1. Recover the DP path (The first 'k' multipliers)
int cur_k = best_sol.k_at_peak;
int cur_idx = best_sol.index_at_peak; // The loop index where we stopped
// We backtrack from cur_idx down to 1
// Note: 'taken' is indexed 1..m.
// 'cur_idx' in the loop `for(int i=0...m)` represents that we have processed up to multipliers[cur_idx-1].
// So we start backtracking from taken[cur_idx].
for (int i = cur_idx; i >= 1; --i) {
if (cur_k == 0) break;
if (taken[i][cur_k]) {
result.push_back(multipliers[i-1].id);
cur_k--;
}
}
reverse(result.begin(), result.end());
// 2. Add "Infinite" remaining multipliers if applicable
if (best_sol.is_infinite) {
// We took 'k' items from 0..cur_idx-1.
// Now we take ALL items from cur_idx..m-1
for (int i = cur_idx; i < m; ++i) {
result.push_back(multipliers[i].id);
}
// And ALL spenders
for (const auto& s : spenders) {
result.push_back(s.id);
}
} else {
// 3. Standard Case: Add Spenders based on budget
// We need the exact money value at best_sol state
// Re-simulate to be safe (or store it). Re-simulation is cheap (60 steps).
long long current_money = (long long)A;
// Apply multipliers in result
for (int id : result) {
// Find component (Inefficient search but O(60*60) is tiny)
// Can map id to details for speed, but vector scan is fine for size 60.
for (const auto& m : multipliers) {
if (m.id == id) {
current_money = (current_money - m.p) * m.t;
if (current_money > INF_CAP) current_money = INF_CAP;
break;
}
}
}
// Add spenders
auto it = upper_bound(spender_prefix.begin(), spender_prefix.end(), current_money);
int count_spenders = (int)(it - spender_prefix.begin()) - 1;
for (int i = 0; i < count_spenders; ++i) {
result.push_back(spenders[i].id);
}
}
return result;
}
