#include <bits/stdc++.h>
using namespace std;
// A is int, P is vector<int>, T is vector<int>
// Return chosen coupon original indices (0-based), matching the Python logic/order.
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int n = (int)P.size();
// ---------- split coupons ----------
vector<pair<long long,int>> cheap; // (price, original_index) for T==1
struct Booster {
long long key; // price * coeff[t]
long long p; // price
int t;
int idx; // original index
};
vector<Booster> boosters;
// coeff = 12*T/(T-1) for T in {2,3,4} -> {24,18,16}
long long coeff[5] = {0, 0, 24, 18, 16};
for (int i = 0; i < n; i++) {
int ti = T[i];
long long pi = (long long)P[i];
if (ti == 1) {
cheap.push_back({pi, i});
} else {
long long key = pi * coeff[ti];
boosters.push_back({key, pi, ti, i});
}
}
// ---------- sort ----------
sort(cheap.begin(), cheap.end(),
[](const auto& a, const auto& b) { return a.first < b.first; });
sort(boosters.begin(), boosters.end(),
[](const Booster& a, const Booster& b) {
if (a.key != b.key) return a.key < b.key;
return a.p < b.p;
});
vector<long long> cheap_psum;
vector<int> cheap_idx;
cheap_psum.reserve(cheap.size() + 1);
cheap_idx.reserve(cheap.size());
cheap_psum.push_back(0);
for (auto &cp : cheap) {
cheap_psum.push_back(cheap_psum.back() + cp.first);
cheap_idx.push_back(cp.second);
}
auto bisect_right = [&](const vector<long long>& arr, long long x) -> int {
return (int)(upper_bound(arr.begin(), arr.end(), x) - arr.begin());
};
// ---------- take all non-decreasing boosters ----------
long long token = (long long)A;
vector<int> prefix_ids;
int i = 0, m = (int)boosters.size();
while (i < m) {
const auto& b = boosters[i];
long long p = b.p;
int t = b.t;
// Python: token * (t - 1) >= p * t
__int128 L = (__int128)token * (t - 1);
__int128 R = (__int128)p * t;
if (L >= R) {
prefix_ids.push_back(b.idx);
token = (token - p) * (long long)t;
i++;
} else {
break;
}
}
// ---------- hard suffix ----------
vector<Booster> hard(boosters.begin() + i, boosters.end());
int h_len = (int)hard.size();
const int K = 61;
vector<long long> dp(K + 1, -1);
dp[0] = token;
vector<vector<uint8_t>> parent(h_len, vector<uint8_t>(K + 1, 0));
for (int j = 0; j < h_len; j++) {
long long p = hard[j].p;
int t = hard[j].t;
int max_c = min(j + 1, K);
for (int c = max_c; c >= 1; c--) {
long long prev = dp[c - 1];
if (prev >= p) {
long long cand = (prev - p) * (long long)t;
if (cand > dp[c]) {
dp[c] = cand;
parent[j][c] = 1;
}
}
}
}
// ---------- evaluate all possibilities ----------
int best_total = -1;
int best_c = 0;
long long best_token = dp[0];
for (int c = 0; c <= min(K, h_len); c++) {
long long cur = dp[c];
if (cur < 0) continue;
int cheap_cnt = bisect_right(cheap_psum, cur) - 1;
int total = i + c + cheap_cnt;
if (total > best_total) {
best_total = total;
best_c = c;
best_token = cur;
}
}
// ---------- reconstruct the chosen hard subset ----------
vector<int> chosen_hard;
int need = best_c;
for (int j = h_len - 1; j >= 0; j--) {
if (need == 0) break;
if (parent[j][need]) {
chosen_hard.push_back(j);
need--;
}
}
reverse(chosen_hard.begin(), chosen_hard.end());
// ---------- build final answer ----------
vector<int> answer;
answer.insert(answer.end(), prefix_ids.begin(), prefix_ids.end());
token = best_token; // as in Python, no need to recompute
for (int j : chosen_hard) {
answer.push_back(hard[j].idx);
}
int cheap_cnt_final = bisect_right(cheap_psum, token) - 1;
for (int k = 0; k < cheap_cnt_final; k++) {
answer.push_back(cheap_idx[k]);
}
return answer;
}
