#include <vector>
#include <algorithm>
#include <iostream>
using namespace std;
typedef long long ll;
struct Coupon {
int id;
ll p;
ll t;
};
// Comparator for Multipliers (T > 1)
// Priority: P=0 first, then by efficiency ratio P*T / (T-1)
bool compareMultipliers(const Coupon& a, const Coupon& b) {
if (a.p == 0 && b.p > 0) return true;
if (b.p == 0 && a.p > 0) return false;
// a.p * a.t / (a.t - 1) < b.p * b.t / (b.t - 1)
// Cross multiply: a.p * a.t * (b.t - 1) < b.p * b.t * (a.t - 1)
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;
}
bool compareSpenders(const Coupon& a, const Coupon& b) {
return a.p < b.p;
}
// Global state tracking for reconstruction
struct BestState {
int total_count; // Total coupons (Multipliers + Spenders)
int dp_k; // How many multipliers were picked via DP
int dp_index; // Index in 'costly_multipliers' where DP stopped
int greedy_taken; // How many *additional* multipliers bought greedily after dp_index
};
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int n = P.size();
vector<Coupon> multipliers;
vector<Coupon> spenders;
// 1. Separate and Sort
for (int i = 0; i < n; ++i) {
if (T[i] == 1) {
spenders.push_back({i, (ll)P[i], (ll)T[i]});
} else {
multipliers.push_back({i, (ll)P[i], (ll)T[i]});
}
}
sort(multipliers.begin(), multipliers.end(), compareMultipliers);
sort(spenders.begin(), spenders.end(), compareSpenders);
// 2. Precompute Spender Costs
int num_spenders = spenders.size();
vector<ll> spender_prefix(num_spenders + 1, 0);
for (int i = 0; i < num_spenders; ++i) {
spender_prefix[i+1] = spender_prefix[i] + spenders[i].p;
}
// 3. Process Free Multipliers (Always take them)
vector<int> result;
ll current_money = A;
// We'll separate multipliers into "Free" (already taken) and "Costly" (for DP)
vector<Coupon> costly_multipliers;
// Safety cap for money (Sum of all costs + buffer)
ll TOTAL_COST_SUM = 0;
for(long long x : P) TOTAL_COST_SUM += x;
ll INF = TOTAL_COST_SUM + 2000;
for (const auto& m : multipliers) {
if (m.p == 0) {
result.push_back(m.id);
if (current_money > 0) {
// Check overflow
if (INF / m.t < current_money) current_money = INF;
else current_money *= m.t;
}
} else {
costly_multipliers.push_back(m);
}
}
// 4. DP Setup
int m = costly_multipliers.size();
// Depth limit: 80 is enough for exponential growth.
// The "Greedy Extension" handles the linear/treadmill cases.
int LIMIT_K = min(m, 80);
// dp[k] = max money after taking k items from costly_multipliers
vector<ll> dp(LIMIT_K + 1, -1);
dp[0] = current_money;
// To reconstruct the DP path:
// parent[i][k] = true if taking item i was the transition to state k
// We use a flat vector approach for memory, storing choices.
// 'choice[i][k]' implies: At step i (processing costly_multipliers[i-1]),
// did we use this item to update dp[k]?
vector<vector<bool>> choice(m + 1, vector<bool>(LIMIT_K + 1, false));
BestState best = {-1, -1, -1, -1};
// Helper: Calculate total coupons given current Money and Multiplier Count
auto check_solution = [&](ll money, int k, int idx, int extra_mul_cnt) {
// How many spenders?
auto it = upper_bound(spender_prefix.begin(), spender_prefix.end(), money);
int cnt_spenders = (int)(it - spender_prefix.begin()) - 1;
// Total = Free (already in result) + DP_k + Greedy_Multipliers + Spenders
// We only track the count relative to the 'costly' set here
int total = (int)result.size() + k + extra_mul_cnt + cnt_spenders;
if (total > best.total_count) {
best = {total, k, idx, extra_mul_cnt};
}
};
// Initial check (0 costly taken)
// Run full greedy simulation from start? No, just check spenders.
// The main loop handles greedy simulations for multipliers.
// But we should check "0 DP items + Greedy Costly" separately?
// The loop runs i from 0 to m. i=0 means before first item.
// We'll integrate the check inside the loop.
// 5. DP Execution
for (int i = 0; i <= m; ++i) {
// A. Greedy Extension Check
// For every valid state dp[k], try to buy remaining multipliers (i...m-1) greedily
for (int k = 0; k <= LIMIT_K; ++k) {
if (dp[k] == -1) continue;
ll sim_money = dp[k];
int extra = 0;
// Optimization: If sim_money >= INF, we buy everything.
if (sim_money >= INF) {
extra = (m - i);
check_solution(sim_money, k, i, extra);
continue;
}
// Run Greedy simulation on remaining costly multipliers
// We limit the lookahead to avoid O(N^2) total time.
// However, with K=80, N=200k, 80*200k = 1.6e7 ops, which is safe for 1-2 sec.
// If TLE occurs, we can optimize this.
bool hit_cap = false;
for (int j = i; j < m; ++j) {
if (sim_money >= costly_multipliers[j].p) {
sim_money = (sim_money - costly_multipliers[j].p) * costly_multipliers[j].t;
extra++;
if (sim_money >= INF) {
sim_money = INF;
// If we hit cap, we can buy all subsequent ones provided P is manageable.
// Ideally we just stop simulation and assume we buy the rest
// IF the rest are "affordable" relative to INF.
// But precise counting is safer. With INF, we likely buy all.
// Let's just bulk add the rest if we are super rich.
extra += (m - 1 - j);
hit_cap = true;
break;
}
}
}
check_solution(sim_money, k, i, extra);
}
if (i == m) break; // Finished checking, no more items to add to DP
// B. Update DP with item i
ll cost = costly_multipliers[i].p;
ll type = costly_multipliers[i].t;
for (int k = LIMIT_K - 1; k >= 0; --k) {
if (dp[k] != -1) {
if (dp[k] >= cost) {
ll next = dp[k] - cost;
if (next > 0) {
if (INF / type < next) next = INF;
else next *= type;
} else {
next = 0;
}
if (next > dp[k+1]) {
dp[k+1] = next;
choice[i+1][k+1] = true;
}
}
}
}
}
// 6. Reconstruction
// A. Add DP Multipliers
// We backtrack from best.dp_index down to 1
vector<int> chosen_ids;
int cur_k = best.dp_k;
// choice is indexed 1..m
// best.dp_index is the step (0..m) where we launched the greedy.
// This means we considered items 0 to dp_index-1 for the DP.
for (int i = best.dp_index; i >= 1; --i) {
if (cur_k == 0) break;
if (choice[i][cur_k]) {
chosen_ids.push_back(costly_multipliers[i-1].id);
cur_k--;
}
}
reverse(chosen_ids.begin(), chosen_ids.end());
result.insert(result.end(), chosen_ids.begin(), chosen_ids.end());
// B. Add Greedy Extension Multipliers
// Recalculate money after DP part
// (We could store it, but recalculation is robust)
ll final_money = A;
// 1. Free ones (already in result)
// 2. DP ones (just added)
// We need to re-simulate exactly to get the money right for the greedy phase
// Since 'result' contains IDs, we need to map them or scan.
// Scanning is fine since result size is small so far.
// Wait, result size is not necessarily small if we had many free ones.
// Better to simulate using the 'costly' array logic.
// Re-simulating logic:
// Money = A.
// Process Free Muls.
// Process DP Muls.
// Optimized Re-sim:
current_money = A;
// Free ones
for(const auto& m : multipliers) {
if(m.p == 0) {
if (current_money > 0) {
if (INF / m.t < current_money) current_money = INF;
else current_money *= m.t;
}
}
}
// DP ones
for(int id : chosen_ids) {
// Find in costly (inefficient search but k<=80)
for(const auto& m : costly_multipliers) {
if(m.id == id) {
current_money = (current_money - m.p) * m.t;
if(current_money > INF) current_money = INF;
break;
}
}
}
// Now run Greedy for the tail
// Starting from best.dp_index
int extra_count = 0;
for (int j = best.dp_index; j < m; ++j) {
// Stop if we satisfied the count found in 'check_solution'
// Actually, we should just run the logic until we can't buy or we hit the count.
// The check_solution calculated 'best.greedy_taken'. We should match that.
if (extra_count == best.greedy_taken) break;
if (current_money >= costly_multipliers[j].p) {
current_money = (current_money - costly_multipliers[j].p) * costly_multipliers[j].t;
if(current_money > INF) current_money = INF;
result.push_back(costly_multipliers[j].id);
extra_count++;
// If infinite, we might have bulk-added in the check.
// If so, we must add all remaining indices.
if (current_money == INF && extra_count < best.greedy_taken) {
// Bulk add remaining
for(int k = j + 1; k < m; ++k) {
if (extra_count == best.greedy_taken) break;
result.push_back(costly_multipliers[k].id);
extra_count++;
}
break;
}
}
}
// C. Add Spenders
auto it = upper_bound(spender_prefix.begin(), spender_prefix.end(), current_money);
int cnt_spenders = (int)(it - spender_prefix.begin()) - 1;
for(int i=0; i<cnt_spenders; ++i) {
result.push_back(spenders[i].id);
}
return result;
}
