#include <bits/stdc++.h>
using namespace std;
struct Coupon {
int idx; // original index
long long price; // P[i]
int type; // T[i] (1..4)
};
struct State {
unsigned long long token; // tokens after this step
int cnt; // number of multi‑type coupons taken so far
int prev; // index of previous state (from previous step)
bool taken; // was this coupon taken at its step?
int step; // step number (0..M)
};
struct Candidate {
unsigned long long token;
int cnt;
int prev;
bool taken;
int step;
};
std::vector<int> max_coupons(int A, std::vector<int> P,
std::vector<int> T) {
const int N = (int)P.size();
// split coupons into multi (type>1) and type1 (type==1)
vector<Coupon> multi, type1;
for (int i = 0; i < N; ++i) {
Coupon c{ i, P[i], T[i] };
if (T[i] == 1) type1.push_back(c);
else multi.push_back(c);
}
// sort multi coupons by increasing w = (price * type) / (type-1)
auto cmpMulti = [](const Coupon& a, const Coupon& b) {
__int128 left = (__int128)a.price * a.type * (b.type - 1);
__int128 right = (__int128)b.price * b.type * (a.type - 1);
if (left != right) return left < right;
return a.idx < b.idx;
};
sort(multi.begin(), multi.end(), cmpMulti);
// sort type1 coupons by price (cheapest first) for later greedy buying
auto cmpType1 = [](const Coupon& a, const Coupon& b) {
if (a.price != b.price) return a.price < b.price;
return a.idx < b.idx;
};
sort(type1.begin(), type1.end(), cmpType1);
// prefix sums of type1 prices (for fast query how many can be bought)
int M1 = (int)type1.size();
vector<unsigned long long> pref(M1 + 1, 0);
for (int i = 0; i < M1; ++i) pref[i + 1] = pref[i] + (unsigned long long)type1[i].price;
// ----- DP over multi coupons -----
const unsigned long long INF = 4'000'000'000'000'000'000ULL; // big enough
vector<State> states;
states.reserve(5000);
states.push_back({ (unsigned long long)A, 0, -1, false, 0 }); // root state
int M = (int)multi.size();
vector<vector<int>> dp(M + 1);
dp[0].push_back(0); // only the root state
for (int i = 1; i <= M; ++i) {
const Coupon& c = multi[i - 1];
vector<Candidate> cand;
cand.reserve(dp[i - 1].size() * 2);
for (int sid : dp[i - 1]) {
const State& st = states[sid];
// skip this coupon
cand.push_back({ st.token, st.cnt, sid, false, i });
// try to take it
if (st.token >= (unsigned long long)c.price) {
unsigned long long diff = st.token - (unsigned long long)c.price;
__int128 prod = (__int128)diff * c.type;
unsigned long long newTok = (prod > INF ? INF : (unsigned long long)prod);
cand.push_back({ newTok, st.cnt + 1, sid, true, i });
}
}
// keep only non‑dominated states (Pareto frontier)
sort(cand.begin(), cand.end(),
[](const Candidate& a, const Candidate& b) {
if (a.token != b.token) return a.token > b.token; // descending token
return a.cnt > b.cnt;
});
vector<int> newIds;
newIds.reserve(cand.size());
int bestCntSoFar = -1;
for (const auto& cc : cand) {
if (cc.cnt > bestCntSoFar) {
State ns;
ns.token = cc.token;
ns.cnt = cc.cnt;
ns.prev = cc.prev;
ns.taken = cc.taken;
ns.step = cc.step;
int nid = (int)states.size();
states.push_back(ns);
newIds.push_back(nid);
bestCntSoFar = cc.cnt;
}
}
dp[i] = std::move(newIds);
}
// ----- choose best final state -----
int bestState = -1;
int bestTotal = -1;
for (int sid : dp[M]) {
const State& st = states[sid];
int canBuyType1 = (int)(upper_bound(pref.begin(), pref.end(), st.token) - pref.begin() - 1);
int total = st.cnt + canBuyType1;
if (total > bestTotal) {
bestTotal = total;
bestState = sid;
}
}
// ----- reconstruct answer -----
vector<int> answer;
// multi coupons (in the order we processed them)
int cur = bestState;
while (states[cur].step > 0) {
if (states[cur].taken) {
int stepIdx = states[cur].step - 1; // index in multi vector
answer.push_back(multi[stepIdx].idx);
}
cur = states[cur].prev;
}
reverse(answer.begin(), answer.end());
// type1 coupons: buy cheapest ones possible
unsigned long long token_end = states[bestState].token;
int k = (int)(upper_bound(pref.begin(), pref.end(), token_end) - pref.begin() - 1);
for (int i = 0; i < k; ++i) answer.push_back(type1[i].idx);
return answer;
}
