#include <bits/stdc++.h>
using namespace std;
using i128 = __int128_t;
struct Coupon {
int idx;
long long P;
int T;
};
static bool cmpMul(const Coupon& a, const Coupon& b) {
// Compare a.P * a.T / (a.T-1) < b.P * b.T / (b.T-1) safely
__int128 left = (__int128)a.P * a.T * (b.T - 1);
__int128 right = (__int128)b.P * b.T * (a.T - 1);
return left < right;
}
static bool cmpOne(const Coupon& a, const Coupon& b) {
return a.P < b.P;
}
// Greedy feasibility: starting with 'tokens', can we take at least 'need' coupons
// from mul[l..M-1] in-order?
static bool can_take_at_least(const vector<Coupon>& mul, int l, int need, i128 tokens) {
const i128 INF = (i128)1 << 120;
int taken = 0;
for (int i = l; i < (int)mul.size() && taken < need; ++i) {
const auto& c = mul[i];
if (tokens >= (i128)c.P) {
tokens = (tokens - (i128)c.P) * (i128)c.T;
if (tokens > INF) tokens = INF;
++taken;
}
}
return taken >= need;
}
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int N = (int)P.size();
vector<Coupon> mul, one;
mul.reserve(N);
one.reserve(N);
for (int i = 0; i < N; ++i) {
Coupon c{ i, (long long)P[i], T[i] };
if (T[i] == 1) one.push_back(c);
else mul.push_back(c);
}
// Sort as before
sort(mul.begin(), mul.end(), cmpMul);
sort(one.begin(), one.end(), cmpOne);
const i128 INF = (i128)1 << 120;
// ---- 1-D DP over number of T>1 picks ----
int M = (int)mul.size();
vector<i128> dp(M + 1, -1), next_dp(M + 1, -1);
dp[0] = (i128)A;
for (int i = 1; i <= M; ++i) {
fill(next_dp.begin(), next_dp.begin() + i + 1, (i128)-1);
const auto& c = mul[i - 1];
// skip
for (int k = 0; k <= i - 1; ++k) if (dp[k] != -1) {
if (next_dp[k] < dp[k]) next_dp[k] = dp[k];
}
// take
for (int k = 1; k <= i; ++k) if (dp[k - 1] != -1) {
i128 prev = dp[k - 1];
if (prev >= (i128)c.P) {
i128 after = (prev - (i128)c.P) * (i128)c.T;
if (after < 0) after = 0;
if (after > INF) after = INF;
if (next_dp[k] < after) next_dp[k] = after;
}
}
dp.swap(next_dp);
}
// ---- choose best k and # of T=1 you can buy afterwards ----
int N1 = (int)one.size();
vector<i128> pref(N1 + 1, 0);
for (int i = 1; i <= N1; ++i) pref[i] = pref[i - 1] + (i128)one[i - 1].P;
auto max_ones = [&](i128 tokens) {
int lo = 0, hi = N1;
while (lo < hi) {
int mid = (lo + hi + 1) >> 1;
if (pref[mid] <= tokens) lo = mid;
else hi = mid - 1;
}
return lo;
};
int bestK = 0, bestOneCnt = 0, bestTotal = 0;
for (int k = 0; k <= M; ++k) if (dp[k] != -1) {
int cntOne = max_ones(dp[k]);
int total = k + cntOne;
if (total > bestTotal) {
bestTotal = total;
bestK = k;
bestOneCnt = cntOne;
}
}
// ---- reconstruct a valid set of multipliers with exactly bestK picks ----
vector<int> orderMul;
orderMul.reserve(bestK);
i128 cur = (i128)A;
int need = bestK;
for (int i = 0; i < M && need > 0; ++i) {
const auto& c = mul[i];
// If we skip this coupon, can we still reach 'need' picks from the suffix?
bool ok_to_skip = can_take_at_least(mul, i + 1, need, cur);
if (!ok_to_skip) {
// Must take this one (feasibility would be lost otherwise).
if (cur < (i128)c.P) {
// Should not happen if DP logic was correct.
// Fall back: skip (keeps algorithm robust).
continue;
}
cur = (cur - (i128)c.P) * (i128)c.T;
if (cur > INF) cur = INF;
orderMul.push_back(c.idx);
--need;
}
// else skip
}
// ---- take first bestOneCnt cheap T=1 coupons ----
vector<int> orderOne;
orderOne.reserve(bestOneCnt);
for (int i = 0; i < bestOneCnt; ++i) orderOne.push_back(one[i].idx);
// ---- concatenate ----
vector<int> res;
res.reserve(bestTotal);
res.insert(res.end(), orderMul.begin(), orderMul.end());
res.insert(res.end(), orderOne.begin(), orderOne.end());
return res;
}
