#include <bits/stdc++.h>
using int64 = long long;
const __int128 INF = (__int128)4e18; // large enough sentinel
struct Coupon {
int64 P; // price
int T; // multiplier (1…4)
int id; // original index
};
/*------------------------------------------------------------*/
/* K-order comparator: renamed to avoid clash with <ratio> */
/*------------------------------------------------------------*/
inline bool coupon_K_less(const Coupon& a, const Coupon& b)
{
// Compare a.P*a.T/(a.T-1) < b.P*b.T/(b.T-1) without division
return (__int128)a.P * a.T * (b.T - 1)
< (__int128)b.P * b.T * (a.T - 1);
}
/*------------------------------------------------------------*/
/* pick optimal subset of multipliers by backwards DP */
/*------------------------------------------------------------*/
struct Node {
int64 need; // minimal starting capital
int cnt; // # multipliers chosen so far
int prev; // predecessor node id (-1 ⇒ root)
int cid; // coupon id used to reach this node (-1 ⇒ root)
};
using NodeId = int;
std::vector<int> choose_multipliers(int64 A, std::vector<Coupon>& mult)
/* returns the ids of chosen multiplier coupons in K-order */
{
std::sort(mult.begin(), mult.end(), coupon_K_less);
/* hash-map: need → node giving largest cnt at that need */
std::unordered_map<int64, NodeId> best;
best.reserve(mult.size() * 2);
std::vector<Node> nodes;
nodes.push_back({0, 0, -1, -1}); // root state
best[0] = 0;
for (int pos = (int)mult.size() - 1; pos >= 0; --pos) {
const auto& c = mult[pos];
std::vector<std::pair<int64,NodeId>> staged; // buffered inserts
for (auto [need, id] : best) {
int64 need2 = c.P + (need + c.T - 1) / c.T; // ceil(need/T)
int cnt2 = nodes[id].cnt + 1;
auto it = best.find(need2);
if (it == best.end() || cnt2 > nodes[it->second].cnt) {
NodeId nid = (NodeId)nodes.size();
nodes.push_back({need2, cnt2, id, c.id});
staged.emplace_back(need2, nid);
}
}
for (auto [k,v] : staged) best[k] = v; // commit
}
/* pick feasible state with maximal cnt */
NodeId v_best = 0;
for (auto [need, id] : best)
if (need <= A && nodes[id].cnt > nodes[v_best].cnt)
v_best = id;
/* reconstruct ids, then reverse to obtain K-order */
std::vector<int> picked;
for (NodeId v = v_best; nodes[v].prev != -1; v = nodes[v].prev)
picked.push_back(nodes[v].cid);
std::reverse(picked.begin(), picked.end());
return picked;
}
/*------------------------------------------------------------*/
/* main wrapper as requested */
/*------------------------------------------------------------*/
std::vector<int> max_coupons(int A,
const std::vector<int>& P_in,
const std::vector<int>& T_in)
{
const int N = (int)P_in.size();
std::vector<Coupon> mult; // T > 1
std::vector<std::pair<int64,int>> flat; // {price,id}
for (int i = 0; i < N; ++i)
(T_in[i] == 1 ? flat : mult).push_back({P_in[i], T_in[i], i});
/* 1️⃣ choose multiplier subset */
std::vector<int> chosen_ids = choose_multipliers(A, mult);
/* gather chosen multipliers in execution order (K-order) */
std::unordered_set<int> chosen_set(chosen_ids.begin(), chosen_ids.end());
std::sort(mult.begin(), mult.end(), coupon_K_less);
std::vector<Coupon> exec;
exec.reserve(chosen_ids.size());
for (const auto& c : mult)
if (chosen_set.count(c.id))
exec.push_back(c);
/* 2️⃣ execute schedule */
__int128 tokens = A; // use 128-bit to avoid overflow
std::vector<int> answer;
answer.reserve(N);
for (const auto& c : exec) {
tokens = (tokens - c.P) * c.T;
if (tokens > INF) tokens = INF;
answer.push_back(c.id);
}
std::sort(flat.begin(), flat.end()); // cheapest flat coupons first
for (auto [price, id] : flat)
if (tokens >= price) {
tokens -= price;
answer.push_back(id);
}
return answer; // 0-indexed buying order
}
