#include <bits/stdc++.h>
using namespace std;
struct ScalingCoupon {
int price;
int mult; // T[i] ( > 1 )
int idx; // original index
};
struct CheapCoupon {
int price;
int idx;
};
static const long long INF = 4000000000000000000LL; // 4·10^18
// comparator: a before b <=> Pa*Ta/(Ta-1) < Pb*Tb/(Tb-1)
bool scalingCmp(const ScalingCoupon& a, const ScalingCoupon& b) {
__int128 left = (__int128)a.price * a.mult * (b.mult - 1);
__int128 right = (__int128)b.price * b.mult * (a.mult - 1);
if (left != right) return left < right;
if (a.price != b.price) return a.price < b.price;
return a.idx < b.idx;
}
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
const int N = (int)P.size();
vector<ScalingCoupon> scaling;
vector<CheapCoupon> cheap;
for (int i = 0; i < N; ++i) {
if (T[i] == 1) {
cheap.push_back({P[i], i});
} else {
scaling.push_back({P[i], T[i], i});
}
}
/* sort scaling coupons by the optimal order (Lemma 3) */
sort(scaling.begin(), scaling.end(), scalingCmp);
/* sort cheap coupons by price, they will be taken greedily */
sort(cheap.begin(), cheap.end(),
[](const CheapCoupon& a, const CheapCoupon& b) {
if (a.price != b.price) return a.price < b.price;
return a.idx < b.idx;
});
int m = (int)scaling.size();
int c1 = (int)cheap.size();
/* DP: dp[i][k] – max tokens after first i scaling coupons,
having taken exactly k of them */
vector<vector<long long>> dp(m + 1, vector<long long>(m + 1, -1));
vector<vector<char>> take(m + 1, vector<char>(m + 1, 0));
dp[0][0] = A;
for (int i = 1; i <= m; ++i) {
const ScalingCoupon& cur = scaling[i - 1];
for (int k = 0; k <= i; ++k) {
// skip current coupon
if (dp[i - 1][k] > dp[i][k]) {
dp[i][k] = dp[i - 1][k];
take[i][k] = 0;
}
// take current coupon
if (k > 0 && dp[i - 1][k - 1] >= cur.price) {
__int128 tmp = (__int128)(dp[i - 1][k - 1] - cur.price) * cur.mult;
long long nxt = (tmp > INF ? INF : (long long)tmp);
if (nxt > dp[i][k]) {
dp[i][k] = nxt;
take[i][k] = 1;
}
}
}
}
/* prefix sums of cheap coupons */
vector<long long> pref(c1 + 1, 0);
for (int i = 0; i < c1; ++i) pref[i + 1] = pref[i] + cheap[i].price;
/* choose the best k */
int bestTotal = -1;
int bestK = 0;
int bestCheapCnt = 0;
long long bestTokens = 0;
for (int k = 0; k <= m; ++k) {
long long tok = dp[m][k];
if (tok < 0) continue;
// maximal cheap coupons affordable
int cheapCnt = (int)(upper_bound(pref.begin(), pref.end(), tok) - pref.begin()) - 1;
int total = k + cheapCnt;
if (total > bestTotal) {
bestTotal = total;
bestK = k;
bestCheapCnt = cheapCnt;
bestTokens = tok;
}
}
/* reconstruct scaling part */
vector<int> answer;
int i = m, k = bestK;
while (i > 0) {
if (take[i][k]) {
answer.push_back(scaling[i - 1].idx);
--k;
}
--i;
}
reverse(answer.begin(), answer.end());
/* add cheap coupons (cheapest first) */
for (int i = 0; i < bestCheapCnt; ++i) {
answer.push_back(cheap[i].idx);
}
return answer;
}
