#include<bits/stdc++.h>
#include "festival.h"
using namespace std;
typedef long long ll;
struct coupon {
int p, t, i;
ll eval(ll x) { return (x-p)*t; }
bool operator<(coupon j) {
if (t==1) return j.p > p;
return j.eval(eval(0)) > eval(j.eval(0));
}
};
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
// exchange argument on coupon order
ll n=P.size(), tot_p=0;
deque<coupon> ord, one;
vector<ll> pref;
for (int i = 0; i < n; i++) {
if (T[i]==1) one.push_back({P[i], 1, i});
else ord.push_back({P[i], T[i], i});
tot_p+=P[i];
} sort(ord.begin(), ord.end());
sort(one.begin(), one.end());
pref.push_back(one[0].p);
for (int i = 1; i < one.size(); i++) {
pref.push_back(pref[i-1]+one[i].p);
}
// cerr << "order\n";
// for (auto i : ord) {
// cerr << i.i << ' ';
// } cerr << '\n';
// for (auto i : one) {
// cerr << i.i << '-' << i.p << ' ';
// } cerr << '\n';
// for (auto i : pref) {
// cerr << i << ' ';
// } cerr << '\n';
// greedy select the coupons that increase tokens taking care of overflow
ll X=A;
vector<int> R;
while (true) {
if (!ord.empty() && X<ord[0].eval(X)) {
R.push_back(ord[0].i);
X=ord[0].eval(X);
tot_p-=ord[0].p;
ord.pop_front(); n--;
} else break;
if (X>=tot_p) {
for (int i = 0; i < n; i++) {
R.push_back(ord[i].i);
} return R;
}
}
// cerr << "greedy\n";
// for (auto i : R) {
// cerr << i << ' ';
// } cerr << '\n';
// for (auto i : ord) {
// cerr << i.i << ' ';
// } cerr << '\n';
// dp[num of coupons bought][pos in coupon sequence]=max possible number of tokens
n=ord.size();
vector<vector<pair<ll, int>>> dp;
dp.push_back(vector<pair<ll, int>>(n+1, {X, -1}));
for (int i = 1; i <= n; i++) {
dp.push_back(vector<pair<ll, int>>(n+1, {-4e18, -1}));
for (int j = i; j <= n; j++) {
int buy=ord[j-1].eval(dp[i-1][j-1].first);
if (buy>dp[i][j-1].first) dp[i][j]={buy, j-1};
else dp[i][j]=dp[i][j-1];
} if (dp[i][n].first<0) break;
}
// cerr << "dp\n";
// for (auto i : dp) {
// for (auto j : i) {
// cerr << j.first << '-' << j.second << ' ';
// } cerr << '\n';
// } cerr << '\n';
// find optimal number of non-ones to buy
int m=dp.size();
// cerr << m << ' ' << n << '\n';
int sz=0, opt=0, j=n;
for (int i = 0; i < m; i++) {
int X=dp[i][n].first;
if (dp[i][n].first<0) break;
int ones=upper_bound(pref.begin(), pref.end(), X)-pref.begin();
// cerr << "in loop: " << ones << ' ' << X << '\n';
if (ones+i>sz) { opt=i; sz=ones+i; }
} sz+=R.size();
// cerr << sz << ' ' << opt << '\n';
// recover sequence
while (opt>0) {
int t=dp[opt][j].second;
R.push_back(ord[t].i);
opt--; j=t-1;
} sz-=R.size();
for (int i = 0; i < sz; i++) {
R.push_back(one[i].i);
}
// cerr << "recovery\n";
return R;
}
