#include <bits/stdc++.h> using namespace std; vector<int> max_coupons(int A, vector<int> P, vector<int> T) { int n = (int)P.size(); vector<pair<int,int>> L1; // (price, idx) for T[i] == 1 struct Coupon { long long p; long long t; int idx; }; vector<Coupon> L2; // (price, t, idx) for T[i] != 1 // Split coupons by type for (int i = 0; i < n; ++i) { if (T[i] == 1) { L1.emplace_back(P[i], i); } else { L2.push_back(Coupon{P[i], T[i], i}); } } // Sort L1 by price sort(L1.begin(), L1.end(), [](const pair<int,int>& a, const pair<int,int>& b) { return a.first < b.first; }); // Build prefix sums for type-1 coupons vector<long long> prefix1; prefix1.reserve(L1.size() + 1); prefix1.push_back(0); for (auto &item : L1) { long long price = item.first; long long next_val = prefix1.back() + price; prefix1.push_back(next_val); } // Custom comparator for L2 (replicates cmp_func) auto cmp_func = [](const Coupon &a, const Coupon &b) { // left = p1 * t1 * (t2 - 1) // right = p2 * t2 * (t1 - 1) __int128 left = (__int128)a.p * a.t * (b.t - 1); __int128 right = (__int128)b.p * b.t * (a.t - 1); if (left != right) { return left < right; } return a.idx < b.idx; }; sort(L2.begin(), L2.end(), cmp_func); int max_k = min(40, (int)L2.size()); const long long NEG_INF = -(long long)1e18; vector<long long> dp(max_k + 1, NEG_INF); vector<vector<int>> seq(max_k + 1); dp[0] = A; seq[0] = {}; // Knapsack-style DP over L2 for (int j = 0; j < (int)L2.size(); ++j) { long long p = L2[j].p; long long t = L2[j].t; int idx = L2[j].idx; for (int taken = max_k; taken >= 1; --taken) { if (dp[taken - 1] >= p) { long long current_token = dp[taken - 1] - p; long long new_token = current_token * t; if (new_token > (long long)1e18) { new_token = (long long)1e18; } if (new_token > dp[taken]) { dp[taken] = new_token; seq[taken] = seq[taken - 1]; seq[taken].push_back(idx); } } } } int best_count = 0; vector<int> best_seq; for (int taken = 0; taken <= max_k; ++taken) { if (dp[taken] < 0) continue; // m = bisect_right(prefix1, dp[taken]) - 1 // prefix1 is sorted non-decreasing int m = (int)(upper_bound(prefix1.begin(), prefix1.end(), dp[taken]) - prefix1.begin()) - 1; int total = taken + m; if (total > best_count) { best_count = total; vector<int> non1_indices = seq[taken]; vector<int> type1_indices; type1_indices.reserve(m); for (int i = 0; i < m; ++i) { type1_indices.push_back(L1[i].second); } best_seq.clear(); best_seq.insert(best_seq.end(), non1_indices.begin(), non1_indices.end()); best_seq.insert(best_seq.end(), type1_indices.begin(), type1_indices.end()); } } return best_seq; }