#include <bits/stdc++.h> using namespace std; std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) { int n = (int)P.size(); // ---------- split coupons ---------- vector<pair<long long,int>> cheap; // (price, original_index) for T=1 struct Booster { long long key; // pi * coeff[ti] int p; int t; int idx; }; vector<Booster> boosters; // for T>1 // coeff = 12*T/(T-1) for T in {2,3,4} unordered_map<int,int> coeff{{2,24},{3,18},{4,16}}; int max_price = 0; for (int i = 0; i < n; i++) { int ti = T[i]; int pi = P[i]; if (ti == 1) { cheap.push_back({pi, i}); } else { long long key = 1LL * pi * coeff[ti]; boosters.push_back({key, pi, ti, i}); max_price = max(max_price, pi); } } // ---------- sort ---------- sort(cheap.begin(), cheap.end(), [](const auto& a, const auto& b){ return a.first < b.first; }); // cheapest first sort(boosters.begin(), boosters.end(), [](const Booster& a, const Booster& b){ if (a.key != b.key) return a.key < b.key; return a.p < b.p; }); vector<long long> cheap_prices; vector<int> cheap_idx; cheap_prices.reserve(cheap.size()); cheap_idx.reserve(cheap.size()); for (auto &pr : cheap) { cheap_prices.push_back(pr.first); cheap_idx.push_back(pr.second); } // prefix sums of cheap prices vector<long long> cheap_psum(cheap_prices.size() + 1, 0); for (size_t i = 0; i < cheap_prices.size(); i++) { cheap_psum[i + 1] = cheap_psum[i] + cheap_prices[i]; } // ---------- take all non-decreasing boosters ---------- long long token = A; vector<int> prefix_ids; // original indices of taken prefix boosters int i = 0; int m = (int)boosters.size(); while (i < m) { const auto &b = boosters[i]; long long p = b.p; long long t = b.t; // token * (t-1) >= p * t if (token * (t - 1) >= p * t) { prefix_ids.push_back(b.idx); token = (token - p) * t; i++; } else { break; } } // ---------- hard suffix ---------- vector<Booster> hard; hard.reserve(m - i); for (int j = i; j < m; j++) hard.push_back(boosters[j]); int h_len = (int)hard.size(); const int K = 61; // safe upper bound vector<long long> dp(K + 1, -1); dp[0] = token; // parent[j][c] = 1 iff hard[j] is used to achieve dp[c] at step j vector<vector<uint8_t>> parent(h_len, vector<uint8_t>(K + 1, 0)); for (int j = 0; j < h_len; j++) { long long p = hard[j].p; long long t = hard[j].t; int max_c = min(j + 1, K); for (int c = max_c; c >= 1; c--) { long long prev = dp[c - 1]; if (prev >= p) { long long cand = (prev - p) * t; if (cand > dp[c]) { dp[c] = cand; parent[j][c] = 1; } } } } // ---------- evaluate all possibilities ---------- int best_total = -1; int best_c = 0; long long best_token = dp[0]; auto cheap_cnt_from_token = [&](long long tok) -> int { // number of cheap items affordable = max k s.t. cheap_psum[k] <= tok auto it = upper_bound(cheap_psum.begin(), cheap_psum.end(), tok); return int(it - cheap_psum.begin()) - 1; }; for (int c = 0; c <= min(K, h_len); c++) { long long cur = dp[c]; if (cur < 0) continue; int cheap_cnt = cheap_cnt_from_token(cur); int total = i + c + cheap_cnt; // i boosters from prefix + c from hard + cheap_cnt if (total > best_total) { best_total = total; best_c = c; best_token = cur; } } // ---------- reconstruct the chosen hard subset ---------- vector<int> chosen_hard_pos; int need = best_c; for (int j = h_len - 1; j >= 0 && need > 0; j--) { if (parent[j][need]) { chosen_hard_pos.push_back(j); need--; } } reverse(chosen_hard_pos.begin(), chosen_hard_pos.end()); // ---------- build the final answer ---------- vector<int> answer; answer.reserve(prefix_ids.size() + chosen_hard_pos.size() + cheap_idx.size()); // prefix boosters for (int id : prefix_ids) answer.push_back(id); // selected hard boosters (append original indices) for (int pos : chosen_hard_pos) { answer.push_back(hard[pos].idx); } // cheap coupons (as many as affordable with best_token) int cheap_cnt_final = cheap_cnt_from_token(best_token); for (int k = 0; k < cheap_cnt_final; k++) { answer.push_back(cheap_idx[k]); } return answer; }