#include <bits/stdc++.h>
using namespace std;
// -------------------------------------------------------------------
// Helper structure for a booster (T > 1)
// -------------------------------------------------------------------
struct Booster {
double a; // sorting key: a = P * T / (T-1)
int price; // P[i]
int type; // T[i] (used later for the token update)
int idx; // original coupon index
};
vector<int> max_coupons(int A,
vector<int> P,
vector<int> T)
{
const int n = static_cast<int>(P.size());
// -----------------------------------------------------------------
// 1) split coupons into boosters (T>1) and cheap coupons (T==1)
// -----------------------------------------------------------------
vector<Booster> boosters;
vector<pair<int,int>> cheap; // (price , original_index)
for (int i = 0; i < n; ++i) {
if (T[i] == 1) {
cheap.emplace_back(P[i], i);
} else {
double a_val = static_cast<double>(P[i]) * T[i] /
(T[i] - 1);
boosters.push_back({a_val, P[i], T[i], i});
}
}
// -----------------------------------------------------------------
// 2) sort the two groups
// -----------------------------------------------------------------
sort(boosters.begin(), boosters.end(),
[](const Booster& x, const Booster& y) { return x.a < y.a; });
sort(cheap.begin(), cheap.end(),
[](const pair<int,int>& x, const pair<int,int>& y) {
return x.first < y.first; // sort by price
});
// cheap_prefix[k] = sum of the k cheapest cheap coupons
// cheap_prefix has length cheap.size() + 1, cheap_prefix[0] == 0
vector<long long> cheap_prefix;
cheap_prefix.reserve(cheap.size() + 1);
cheap_prefix.push_back(0);
for (auto &c : cheap) {
cheap_prefix.push_back(cheap_prefix.back() + static_cast<long long>(c.first));
}
// -----------------------------------------------------------------
// 3) DP over boosters, keeping only non‑dominated states
// -----------------------------------------------------------------
// DP vectors – they are indexed by the *state id* (its position in the vectors)
vector<long long> tokens; // tokens after the state
vector<int> cnts; // number of boosters taken to reach the state
vector<int> prevs; // previous state id ( -1 for the start state )
vector<int> idxs; // original coupon index that created the state
// start state (state id = 0)
tokens.push_back(A); // we start with A tokens
cnts.push_back(0);
prevs.push_back(-1);
idxs.push_back(-1); // sentinel
// frontier = list of state ids that are currently non‑dominated
vector<int> frontier = {0};
// iterate over boosters in increasing a‑value
for (const Booster& B : boosters) {
vector<int> new_state_ids;
// try to take this booster from every state that is currently reachable
for (int sid : frontier) {
if (tokens[sid] >= B.price) { // we can afford the price
long long new_tok = (tokens[sid] - B.price) * B.type;
// create a new DP state
tokens.push_back(new_tok);
cnts.push_back(cnts[sid] + 1);
prevs.push_back(sid);
idxs.push_back(B.idx);
new_state_ids.push_back(static_cast<int>(tokens.size()) - 1);
}
}
// -------------------------------------------------------------
// Merge old frontier and new states, then keep only non‑dominated
// -------------------------------------------------------------
vector<int> candidates = frontier;
candidates.insert(candidates.end(), new_state_ids.begin(),
new_state_ids.end());
// sort by decreasing tokens, then decreasing cnt (exactly as in Python)
sort(candidates.begin(), candidates.end(),
[&](int s1, int s2) {
if (tokens[s1] != tokens[s2]) return tokens[s1] > tokens[s2];
return cnts[s1] > cnts[s2];
});
// keep states that are *strictly* better in cnt than any previously kept
vector<int> new_frontier;
int best_cnt = -1;
for (int sid : candidates) {
if (cnts[sid] > best_cnt) {
new_frontier.push_back(sid);
best_cnt = cnts[sid];
}
}
frontier.swap(new_frontier);
}
// -----------------------------------------------------------------
// 4) Choose the best final state (booster + cheap coupons)
// -----------------------------------------------------------------
int best_state = -1;
int best_cheap = 0; // number of cheap coupons we will finally take
int best_total = -1; // cnt(booster) + cheap
for (int sid : frontier) {
int st_cnt = cnts[sid];
long long st_tok = tokens[sid];
// how many cheap coupons can we still afford?
// cheap_prefix is monotone, we need the largest k with cheap_prefix[k] <= st_tok
int cheap_m = static_cast<int>(
upper_bound(cheap_prefix.begin(), cheap_prefix.end(), st_tok) - cheap_prefix.begin()
) - 1; // -1 because prefix array has a leading 0
int total = st_cnt + cheap_m;
if (total > best_total) {
best_total = total;
best_state = sid;
best_cheap = cheap_m;
}
}
// -----------------------------------------------------------------
// 5) Reconstruct the answer (order of taken coupons)
// -----------------------------------------------------------------
vector<int> answer;
// booster part (trace back the DP chain)
for (int sid = best_state; sid != 0; sid = prevs[sid]) {
answer.push_back(idxs[sid]); // original index of this booster
}
reverse(answer.begin(), answer.end()); // chronological order
// cheap coupons – they were already sorted by price
for (int i = 0; i < best_cheap; ++i) {
answer.push_back(cheap[i].second); // original index of cheap coupon i
}
return answer;
}
