#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
struct Coupon { ll price; int mult, idx; };
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int n = P.size();
vector<Coupon> S;
vector<pair<ll,int>> L;
// Separate coupons
for(int i = 0; i < n; i++) {
if(T[i] >= 2) {
S.push_back({P[i], T[i], i});
} else {
L.emplace_back(P[i], i);
}
}
sort(L.begin(), L.end()); // type-1 sorted by price
int m = S.size();
// dp2[i][j] = max tokens after considering first i S-coupons and picking j of them
const ll NEG = -1;
vector<vector<ll>> dp2(m+1, vector<ll>(m+1, NEG));
struct Pre { int prev_j; bool took; };
vector<vector<Pre>> pre(m+1, vector<Pre>(m+1, {-1,false}));
dp2[0][0] = A;
for(int i = 0; i < m; i++) {
ll price = S[i].price;
int mult = S[i].mult;
for(int j = 0; j <= i; j++) {
if(dp2[i][j] < 0) continue;
// skip S[i]
if(dp2[i][j] > dp2[i+1][j]) {
dp2[i+1][j] = dp2[i][j];
pre[i+1][j] = {j, false};
}
// take S[i]
if(dp2[i][j] >= price) {
ll nv = (dp2[i][j] - price) * mult;
if(nv > dp2[i+1][j+1]) {
dp2[i+1][j+1] = nv;
pre[i+1][j+1] = {j, true};
}
}
}
}
// find best total count
int bestj = 0;
int bestCount = 0;
ll bestTokens = A;
for(int j = 0; j <= m; j++) {
ll tokens = dp2[m][j];
if(tokens < 0) continue;
// count type-1 buys
ll rem = tokens;
int cnt1 = 0;
for(auto &pr : L) {
if(pr.first <= rem) { rem -= pr.first; cnt1++; }
else break;
}
if(j + cnt1 > bestCount) {
bestCount = j + cnt1;
bestj = j;
bestTokens = tokens;
}
}
// backtrack S-coupons
vector<int> result;
int ci = m, cj = bestj;
while(ci > 0) {
Pre p = pre[ci][cj];
if(p.prev_j < 0) break;
if(p.took) {
// took S[ci-1]
result.push_back(S[ci-1].idx);
}
cj = p.prev_j;
ci--;
}
reverse(result.begin(), result.end());
// append type-1 coupons
ll rem = bestTokens;
for(auto &pr : L) {
if(pr.first <= rem) {
rem -= pr.first;
result.push_back(pr.second);
} else break;
}
return result;
}
