#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
struct Coupon {
int idx;
long long p;
int t;
// used only for sorting
};
static bool cmpCoupon(const Coupon& a, const Coupon& b) {
// T == 1 has infinite ratio, should be placed at the end
if (a.t == 1 && b.t == 1) return a.idx < b.idx;
if (a.t == 1) return false; // a after b
if (b.t == 1) return true; // a before b
// compare a.p * a.t / (a.t-1) < b.p * b.t / (b.t-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;
}
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
const long long INF = 4000000000000000000LL; // 4e18, safely fits in int64
int N = (int)P.size();
vector<Coupon> coupons(N);
for (int i = 0; i < N; ++i) {
coupons[i] = {i, (long long)P[i], T[i]};
}
sort(coupons.begin(), coupons.end(), cmpCoupon);
// ----- phase 1 : take all "good" coupons (cur >= R_i) -----
__int128 cur = A;
vector<int> answer;
int splitPos = N; // first index that is not taken in the first phase
for (int i = 0; i < N; ++i) {
const Coupon& c = coupons[i];
bool good = false;
if (c.t != 1) {
__int128 left = cur * (c.t - 1);
__int128 right = (__int128)c.p * c.t;
if (left >= right) good = true; // cur >= R_i
}
// (for T==1, good stays false)
if (good) {
// affordable is guaranteed because cur >= R_i > p
answer.push_back(c.idx);
cur = (cur - c.p) * c.t;
if (cur > (__int128)INF) cur = (__int128)INF;
} else {
splitPos = i;
break;
}
}
// ----- collect remaining coupons -----
vector<Coupon> highT; // T >= 2, still lossful
vector<pair<long long,int>> type1; // (price, idx) for T==1
for (int i = splitPos; i < N; ++i) {
const Coupon& c = coupons[i];
if (c.t == 1) {
type1.emplace_back(c.p, c.idx);
} else {
highT.push_back(c);
}
}
// sort type1 by price (cheapest first)
sort(type1.begin(), type1.end(),
[](const pair<long long,int>& a, const pair<long long,int>& b){
if (a.first != b.first) return a.first < b.first;
return a.second < b.second;
});
// prefix sums of type1 prices
int M1 = (int)type1.size();
vector<long long> pref(M1 + 1, 0);
for (int i = 0; i < M1; ++i) {
pref[i+1] = pref[i] + type1[i].first;
}
// ----- DP over lossful highT coupons -----
const int KMAX = 70; // enough for any feasible number
int M = (int)highT.size();
int maxK = min(KMAX, M);
// dp[i][k] = maximal remaining tokens after processing first i highT coupons,
// having taken exactly k of them ( -1 means impossible )
vector<vector<long long>> dp(M+1, vector<long long>(maxK+2, -1));
vector<vector<char>> take(M+1, vector<char>(maxK+2, 0)); // 1 if taken this coupon
long long cur0 = (cur > (__int128)INF) ? INF : (long long)cur;
dp[0][0] = cur0;
for (int i = 0; i < M; ++i) {
const Coupon& c = highT[i];
for (int k = 0; k <= maxK; ++k) {
if (dp[i][k] < 0) continue;
// skip this coupon
if (dp[i][k] > dp[i+1][k]) {
dp[i+1][k] = dp[i][k];
take[i+1][k] = 0;
}
// take this coupon if affordable
if (k+1 <= maxK && dp[i][k] >= c.p) {
__int128 nxt = (__int128)(dp[i][k] - c.p) * c.t;
long long nxt_ll = (nxt > (__int128)INF) ? INF : (long long)nxt;
if (nxt_ll > dp[i+1][k+1]) {
dp[i+1][k+1] = nxt_ll;
take[i+1][k+1] = 1; // came from taking coupon i
}
}
}
}
// ----- decide best combination of highT and type1 -----
int bestHighCnt = 0;
int bestType1Cnt = 0;
long long bestRemain = 0;
int limitK = min(maxK, M);
for (int k = 0; k <= limitK; ++k) {
if (dp[M][k] < 0) continue;
long long remain = dp[M][k];
// how many type1 coupons can we afford?
int cnt1 = (int)(upper_bound(pref.begin(), pref.end(), remain) - pref.begin()) - 1;
if (cnt1 < 0) cnt1 = 0;
int total = (int)answer.size() + k + cnt1;
if (total > (int)answer.size() + bestHighCnt + bestType1Cnt) {
bestHighCnt = k;
bestType1Cnt = cnt1;
bestRemain = remain;
}
}
// ----- reconstruct selected highT coupons -----
vector<int> selectedHigh;
if (bestHighCnt > 0) {
int i = M;
int k = bestHighCnt;
while (i > 0) {
if (take[i][k]) {
selectedHigh.push_back(highT[i-1].idx);
--k;
}
--i;
}
reverse(selectedHigh.begin(), selectedHigh.end());
}
// ----- select type1 coupons -----
vector<int> selectedType1;
for (int i = 0; i < bestType1Cnt; ++i) {
selectedType1.push_back(type1[i].second);
}
// ----- final answer -----
// answer already contains good coupons (in correct order)
answer.insert(answer.end(), selectedHigh.begin(), selectedHigh.end());
answer.insert(answer.end(), selectedType1.begin(), selectedType1.end());
return answer;
}
