#include <bits/stdc++.h>
using namespace std;
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int n = (int)P.size();
// L1: (price, index) for T[i] == 1
vector<pair<long long,int>> L1;
// L2: coupons with t > 1
struct Coupon {
long long p;
long long t;
int idx;
};
vector<Coupon> L2;
for (int i = 0; i < n; ++i) {
if (T[i] == 1) {
L1.push_back({(long long)P[i], i});
} else {
L2.push_back({(long long)P[i], (long long)T[i], i});
}
}
// Sort L1 by price
sort(L1.begin(), L1.end(),
[](const pair<long long,int> &a, const pair<long long,int> &b) {
return a.first < b.first;
});
// Sort L2 with custom comparator
if (!L2.empty()) {
sort(L2.begin(), L2.end(),
[](const Coupon &a, const Coupon &b) {
// Use __int128 to avoid overflow
__int128 p1 = a.p, t1 = a.t;
__int128 p2 = b.p, t2 = b.t;
__int128 left = p1 * t1 * (t2 - 1);
__int128 right = p2 * t2 * (t1 - 1);
return left < right;
});
}
// Prefix sums for L1 prices
vector<long long> prefix1;
prefix1.reserve(L1.size() + 1);
prefix1.push_back(0);
for (auto &pr : L1) {
long long price = pr.first;
long long next_val = prefix1.back() + price;
prefix1.push_back(next_val);
}
const int max_set2 = 100;
const long long CLAMP_MAX = (long long)1e18;
// dp[k] = max money achievable after using k coupons from L2
vector<long long> dp(max_set2 + 1, -1); // -1 means unreachable
dp[0] = (long long)A;
// parent[k] = (previous k, original index of coupon)
vector<pair<int,int>> parent(max_set2 + 1, {-1, -1});
// Knapsack over L2 coupons
for (const auto &coupon : L2) {
long long p = coupon.p;
long long t = coupon.t;
int orig_i = coupon.idx;
for (int k = max_set2 - 1; k >= 0; --k) {
if (dp[k] < p || dp[k] < 0) continue;
long long current_val = dp[k];
__int128 new_val_128 = (__int128)(current_val - p) * t;
long long new_val;
if (new_val_128 > CLAMP_MAX) new_val = CLAMP_MAX;
else new_val = (long long)new_val_128;
if (new_val > dp[k + 1]) {
dp[k + 1] = new_val;
parent[k + 1] = {k, orig_i};
}
}
}
// Choose best combination of L2 coupons (k2) and L1 items (m)
int best_total = 0;
int best_k2 = 0;
int best_m = 0;
for (int k2 = 0; k2 <= max_set2; ++k2) {
if (dp[k2] < 0) continue;
long long money = dp[k2];
// m = number of items from L1 we can buy with 'money'
int m = (int)(upper_bound(prefix1.begin(), prefix1.end(), money) - prefix1.begin()) - 1;
int total = k2 + m;
if (total > best_total) {
best_total = total;
best_k2 = k2;
best_m = m;
}
}
// Reconstruct which L2 coupons were used
vector<int> seq_set2;
{
int k = best_k2;
while (k > 0) {
auto [prev_k, idx] = parent[k];
seq_set2.push_back(idx);
k = prev_k;
}
reverse(seq_set2.begin(), seq_set2.end());
}
// Take first best_m items from L1 (they are cheapest due to sorting)
vector<int> seq_set1;
for (int i = 0; i < best_m; ++i) {
seq_set1.push_back(L1[i].second);
}
// Return concatenation of seq_set2 + seq_set1
vector<int> result;
result.reserve(seq_set2.size() + seq_set1.size());
result.insert(result.end(), seq_set2.begin(), seq_set2.end());
result.insert(result.end(), seq_set1.begin(), seq_set1.end());
return result;
}
