#include <vector>
#include <algorithm>
#include <iostream>
#include <queue>
#include <tuple>
using namespace std;
// Use explicit long long to prevent overflow
typedef long long ll;
// State structure for BFS queue (compact to save memory/time)
struct State {
uint8_t u1; // Count of Type 2 coupons bought in BFS phase
uint8_t u2; // Count of Type 3 coupons bought in BFS phase
uint8_t u3; // Count of Type 4 coupons bought in BFS phase
};
// DP Array to replace Map
// Dimensions: [65][65][65]. 65 is safe because 2^60 > 10^18.
// Stores the max money achievable for a specific count of T2, T3, T4 coupons.
ll dp[65][65][65];
struct Candidate {
int t_idx; // 1 for T2, 2 for T3, 3 for T4
int v_idx; // index in the respective sorted vector
ll price;
};
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
int n = P.size();
// 1. Separate coupons into sorted buckets
// v[0]=Spenders, v[1]=Type2, v[2]=Type3, v[3]=Type4
vector<vector<pair<ll, int>>> v(4);
ll sum_all_prices = 0;
for(int i=0; i<n; ++i) {
if(T[i] >= 1 && T[i] <= 4) {
v[T[i]-1].push_back({(ll)P[i], i});
sum_all_prices += P[i];
}
}
// Sort buckets by Price (Cheapest first)
for(int i=0; i<4; ++i) sort(v[i].begin(), v[i].end());
// Define "Infinite Money" Cap
// If money exceeds this, we can buy everything.
// Use the larger of a safe baseline or the actual sum of prices.
ll MAX_VAL = 2000000000000000LL; // 2e15 baseline
if (sum_all_prices + 2000 > MAX_VAL) MAX_VAL = sum_all_prices + 2000;
// Cap to avoid LL overflow during multiplications
if (MAX_VAL > 4000000000000000000LL) MAX_VAL = 4000000000000000000LL;
// 2. Profit Phase Candidates (Heuristic)
vector<Candidate> candidates;
for(int t=1; t<4; ++t) {
for(int i=0; i<v[t].size(); ++i) {
candidates.push_back({t, i, v[t][i].first});
}
}
// Sort by efficiency P*T / (T-1) using cross-multiplication
sort(candidates.begin(), candidates.end(), [](const Candidate& a, const Candidate& b){
ll Ta = a.t_idx + 1;
ll Tb = b.t_idx + 1;
unsigned __int128 v1 = (unsigned __int128)a.price * Ta * (Tb - 1);
unsigned __int128 v2 = (unsigned __int128)b.price * Tb * (Ta - 1);
return v1 < v2;
});
// 3. Execute Profit Phase
// Greedily buy anything that strictly increases our money.
vector<int> res_indices;
int c[4] = {0, 0, 0, 0}; // c[1]=count of T2 used, etc.
ll cur = A;
for(const auto& cand : candidates) {
int t = cand.t_idx;
if(cand.v_idx == c[t]) {
ll p = v[t][c[t]].first;
if(cur < p) continue;
ll mult = t + 1;
ll nxt = (cur - p) * mult;
if(nxt > MAX_VAL) nxt = MAX_VAL;
if(nxt > cur) {
cur = nxt;
res_indices.push_back(v[t][c[t]].second);
c[t]++;
}
}
}
// 4. BFS Phase Setup
int base[4]; // Offset from profit phase
base[1] = c[1]; base[2] = c[2]; base[3] = c[3];
// Reset DP Array to -1
for(int i=0; i<65; ++i)
for(int j=0; j<65; ++j)
for(int k=0; k<65; ++k)
dp[i][j][k] = -1;
dp[0][0][0] = cur;
queue<State> q;
q.push({0,0,0});
int max_total = 0;
State best_state = {0,0,0};
// Precompute Spender (T1) prefix sums for fast lookup
vector<ll> p1;
p1.push_back(0);
for(auto& x : v[0]) p1.push_back(p1.back() + x.first);
// Helper to update global best
auto update = [&](ll money, int u1, int u2, int u3) {
auto it = upper_bound(p1.begin(), p1.end(), money);
int c1 = (int)(it - p1.begin()) - 1;
int tot = (int)res_indices.size() + u1 + u2 + u3 + c1;
if(tot > max_total) {
max_total = tot;
best_state = {(uint8_t)u1, (uint8_t)u2, (uint8_t)u3};
}
};
update(cur, 0, 0, 0);
// 5. BFS Execution
while(!q.empty()) {
State s = q.front();
q.pop();
int u1 = s.u1, u2 = s.u2, u3 = s.u3;
ll m = dp[u1][u2][u3];
// PRUNING: If money is saturated, we can buy EVERYTHING remaining.
if(m >= MAX_VAL) {
// Calculate remaining items for T2, T3, T4
int rem2 = max(0, (int)v[1].size() - (base[1] + u1));
int rem3 = max(0, (int)v[2].size() - (base[2] + u2));
int rem4 = max(0, (int)v[3].size() - (base[3] + u3));
int rem1 = (int)v[0].size(); // All spenders
int tot = (int)res_indices.size() + u1 + u2 + u3 + rem1 + rem2 + rem3 + rem4;
if(tot > max_total) {
max_total = tot;
best_state = s;
}
continue; // Stop expanding this saturated branch
}
// Try adding next Type 2
if(base[1] + u1 < v[1].size() && u1 + 1 < 65) {
ll p = v[1][base[1] + u1].first;
if(m >= p) {
ll nm = (m - p) * 2;
if(nm > MAX_VAL) nm = MAX_VAL;
if(nm > dp[u1+1][u2][u3]) {
dp[u1+1][u2][u3] = nm;
q.push({(uint8_t)(u1+1), (uint8_t)u2, (uint8_t)u3});
update(nm, u1+1, u2, u3);
}
}
}
// Try adding next Type 3
if(base[2] + u2 < v[2].size() && u2 + 1 < 65) {
ll p = v[2][base[2] + u2].first;
if(m >= p) {
ll nm = (m - p) * 3;
if(nm > MAX_VAL) nm = MAX_VAL;
if(nm > dp[u1][u2+1][u3]) {
dp[u1][u2+1][u3] = nm;
q.push({(uint8_t)u1, (uint8_t)(u2+1), (uint8_t)u3});
update(nm, u1, u2+1, u3);
}
}
}
// Try adding next Type 4
if(base[3] + u3 < v[3].size() && u3 + 1 < 65) {
ll p = v[3][base[3] + u3].first;
if(m >= p) {
ll nm = (m - p) * 4;
if(nm > MAX_VAL) nm = MAX_VAL;
if(nm > dp[u1][u2][u3+1]) {
dp[u1][u2][u3+1] = nm;
q.push({(uint8_t)u1, (uint8_t)u2, (uint8_t)(u3+1)});
update(nm, u1, u2, u3+1);
}
}
}
}
// 6. Reconstruction
// Backtrack from best_state to find the BFS coupons
vector<int> bfs_res;
int cu1 = best_state.u1;
int cu2 = best_state.u2;
int cu3 = best_state.u3;
while(cu1 > 0 || cu2 > 0 || cu3 > 0) {
ll curr_m = dp[cu1][cu2][cu3];
bool found = false;
// Try T2 parent
if(cu1 > 0) {
ll prev = dp[cu1-1][cu2][cu3];
if(prev != -1) {
ll p = v[1][base[1] + cu1 - 1].first;
ll check = (prev - p) * 2;
if(check > MAX_VAL) check = MAX_VAL;
if(check == curr_m) {
bfs_res.push_back(v[1][base[1] + cu1 - 1].second);
cu1--; found = true;
}
}
}
// Try T3 parent
if(!found && cu2 > 0) {
ll prev = dp[cu1][cu2-1][cu3];
if(prev != -1) {
ll p = v[2][base[2] + cu2 - 1].first;
ll check = (prev - p) * 3;
if(check > MAX_VAL) check = MAX_VAL;
if(check == curr_m) {
bfs_res.push_back(v[2][base[2] + cu2 - 1].second);
cu2--; found = true;
}
}
}
// Try T4 parent
if(!found && cu3 > 0) {
ll prev = dp[cu1][cu2][cu3-1];
if(prev != -1) {
ll p = v[3][base[3] + cu3 - 1].first;
ll check = (prev - p) * 4;
if(check > MAX_VAL) check = MAX_VAL;
if(check == curr_m) {
bfs_res.push_back(v[3][base[3] + cu3 - 1].second);
cu3--; found = true;
}
}
}
if(!found) break; // Should not happen
}
reverse(bfs_res.begin(), bfs_res.end());
res_indices.insert(res_indices.end(), bfs_res.begin(), bfs_res.end());
// 7. Add Remaining Items
ll fin_m = dp[best_state.u1][best_state.u2][best_state.u3];
// Add Spenders (Type 1)
auto it = upper_bound(p1.begin(), p1.end(), fin_m);
int c1 = (int)(it - p1.begin()) - 1;
for(int i=0; i<c1; ++i) res_indices.push_back(v[0][i].second);
// If we hit Infinite Money, add all remaining Multipliers too
if(fin_m >= MAX_VAL) {
// T2
for(size_t i = base[1] + best_state.u1; i < v[1].size(); ++i)
res_indices.push_back(v[1][i].second);
// T3
for(size_t i = base[2] + best_state.u2; i < v[2].size(); ++i)
res_indices.push_back(v[2][i].second);
// T4
for(size_t i = base[3] + best_state.u3; i < v[3].size(); ++i)
res_indices.push_back(v[3][i].second);
}
return res_indices;
}
