#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;
}
