#include <vector>
#include <algorithm>
#include <iostream>
#include <map>
#include <queue>
using namespace std;
// Use explicit long long for money to prevent overflow
typedef long long ll;
const ll MAX_VAL = 1e16; // Cap money at 10^16 (enough to buy everything)
// Structure to manage coupons in the sorting phase
struct Candidate {
int type_idx; // 1, 2, or 3 representing Types 2, 3, 4
int vec_idx; // Index within the specific type vector
ll price;
};
// Standard function signature required by the problem
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
int n = P.size();
// 1. Separate coupons into buckets based on Type
// v[0] stores Type 1 (Spenders)
// v[1] stores Type 2
// v[2] stores Type 3
// v[3] stores Type 4
vector<vector<pair<ll, int>>> v(4);
for (int i = 0; i < n; ++i) {
if (T[i] >= 1 && T[i] <= 4) {
v[T[i] - 1].push_back({(ll)P[i], i});
}
}
// 2. Sort each bucket by Price (Cheapest first)
for (int i = 0; i < 4; ++i) {
sort(v[i].begin(), v[i].end());
}
// 3. Prepare "Profit Phase" Candidates (Types 2, 3, 4)
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 candidates by Efficiency: P*T / (T-1)
// We use cross-multiplication to compare: A < B
// Formula: Pa * Ta * (Tb - 1) < Pb * Tb * (Ta - 1)
sort(candidates.begin(), candidates.end(), [&](const Candidate& a, const Candidate& b) {
ll Ta = a.type_idx + 1;
ll Tb = b.type_idx + 1;
// Use __int128 to ensure no overflow during comparison
unsigned __int128 val1 = (unsigned __int128)a.price * Ta * (Tb - 1);
unsigned __int128 val2 = (unsigned __int128)b.price * Tb * (Ta - 1);
return val1 < val2;
});
// 4. Execute "Profit Phase" (Greedy)
// Buy any coupon that results in strictly MORE money than before.
vector<int> result_indices;
vector<int> counters(4, 0); // Tracks how many of each type we have used
ll current_money = A;
for (const auto& c : candidates) {
int t = c.type_idx;
int idx = c.vec_idx;
// We can only process this candidate if it is the NEXT available one in its bucket.
// Due to efficiency sorting, valid candidates usually appear in order,
// but we must skip "future" expensive ones if we haven't bought cheaper ones yet.
if (idx == counters[t]) {
ll price = v[t][idx].first;
if (current_money < price) {
// Cannot afford this specific coupon.
// Since candidates are sorted by efficiency (easier affordability),
// we likely can't afford subsequent ones either, but we continue to be safe.
continue;
}
ll mult = t + 1;
ll next_val = (current_money - price) * mult;
// Cap the value
if (next_val > MAX_VAL) next_val = MAX_VAL;
if (next_val > current_money) {
current_money = next_val;
result_indices.push_back(v[t][idx].second);
counters[t]++;
} else {
// If the most efficient available coupon is not profitable,
// no other coupon will be. We stop the Profit Phase.
break;
}
}
}
// 5. BFS Phase (State Space Search)
// We explore buying "Costly" coupons (where Money decreases or stays same).
// State: {count_type2, count_type3, count_type4} (Offsets from 'counters')
// Base offsets from profit phase
int base_c1 = counters[1];
int base_c2 = counters[2];
int base_c3 = counters[3];
// DP Map: State Vector -> Max Money
map<vector<int>, ll> dp;
// Parent Map for reconstruction: State Vector -> {TypeAdded, IndexInResultVector}
// TypeAdded: 1 for T2, 2 for T3, 3 for T4
map<vector<int>, pair<int, int>> parent;
vector<int> start_state = {0, 0, 0};
dp[start_state] = current_money;
queue<vector<int>> q;
q.push(start_state);
// Track the global maximum coupons found
int max_coupons_found = 0;
vector<int> best_state_vector = start_state;
// Precompute Spender (Type 1) prefix sums for O(log N) lookup
vector<ll> type1_prefix;
type1_prefix.push_back(0);
for (const auto& p : v[0]) {
type1_prefix.push_back(type1_prefix.back() + p.first);
}
// Helper to check total coupons (Multipliers + Spenders)
auto check_solution = [&](ll money, int extra_mults, const vector<int>& state) {
auto it = upper_bound(type1_prefix.begin(), type1_prefix.end(), money);
int spenders_count = (int)(it - type1_prefix.begin()) - 1;
int total = (int)result_indices.size() + extra_mults + spenders_count;
if (total > max_coupons_found) {
max_coupons_found = total;
best_state_vector = state;
}
};
// Initial check
check_solution(current_money, 0, start_state);
while (!q.empty()) {
vector<int> s = q.front();
q.pop();
ll cur_m = dp[s];
int u1 = s[0];
int u2 = s[1];
int u3 = s[2];
int total_extra = u1 + u2 + u3;
// Try adding Type 2 (Index 1 in v)
if (base_c1 + u1 < v[1].size()) {
ll p = v[1][base_c1 + u1].first;
if (cur_m >= p) {
ll next_m = (cur_m - p) * 2;
if (next_m > MAX_VAL) next_m = MAX_VAL;
vector<int> next_s = {u1 + 1, u2, u3};
if (dp.find(next_s) == dp.end() || dp[next_s] < next_m) {
dp[next_s] = next_m;
parent[next_s] = {1, v[1][base_c1 + u1].second};
q.push(next_s);
check_solution(next_m, total_extra + 1, next_s);
}
}
}
// Try adding Type 3 (Index 2 in v)
if (base_c2 + u2 < v[2].size()) {
ll p = v[2][base_c2 + u2].first;
if (cur_m >= p) {
ll next_m = (cur_m - p) * 3;
if (next_m > MAX_VAL) next_m = MAX_VAL;
vector<int> next_s = {u1, u2 + 1, u3};
if (dp.find(next_s) == dp.end() || dp[next_s] < next_m) {
dp[next_s] = next_m;
parent[next_s] = {2, v[2][base_c2 + u2].second};
q.push(next_s);
check_solution(next_m, total_extra + 1, next_s);
}
}
}
// Try adding Type 4 (Index 3 in v)
if (base_c3 + u3 < v[3].size()) {
ll p = v[3][base_c3 + u3].first;
if (cur_m >= p) {
ll next_m = (cur_m - p) * 4;
if (next_m > MAX_VAL) next_m = MAX_VAL;
vector<int> next_s = {u1, u2, u3 + 1};
if (dp.find(next_s) == dp.end() || dp[next_s] < next_m) {
dp[next_s] = next_m;
parent[next_s] = {3, v[3][base_c3 + u3].second};
q.push(next_s);
check_solution(next_m, total_extra + 1, next_s);
}
}
}
}
// 6. Reconstruction
// Retrieve the BFS multipliers by backtracking from best_state_vector
vector<int> bfs_indices;
vector<int> curr = best_state_vector;
while (curr[0] > 0 || curr[1] > 0 || curr[2] > 0) {
// Safety check to prevent RTE if state is missing (should not happen)
if (parent.find(curr) == parent.end()) break;
pair<int, int> info = parent[curr];
bfs_indices.push_back(info.second);
// Backtrack state: info.first is 1, 2, or 3 corresponding to index 0, 1, 2
curr[info.first - 1]--;
}
// BFS reconstruction is in reverse order (last bought first)
reverse(bfs_indices.begin(), bfs_indices.end());
// Append BFS results to the Profit Phase results
result_indices.insert(result_indices.end(), bfs_indices.begin(), bfs_indices.end());
// 7. Add Spenders
// Calculate final money based on the optimal DP state
ll final_money = dp[best_state_vector];
auto it = upper_bound(type1_prefix.begin(), type1_prefix.end(), final_money);
int cnt_spenders = (int)(it - type1_prefix.begin()) - 1;
for (int i = 0; i < cnt_spenders; ++i) {
result_indices.push_back(v[0][i].second);
}
return result_indices;
}
