#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
#define all(x) (x).begin(), (x).end()
typedef long long ll;
using pi = pair<ll, ll>;
using vi = vector<ll>;
struct coupon {
ll t, p, idx;
ll eval(ll R) {
ll ans = (R-p)*t;
if (R < p) ans = -4e18;
else if (R >= 1e16) ans = R;
return ans;
}
bool norm(ll x) {
return (t * -p) <= (1-t) * x;
}
const bool operator<(const coupon &other) { // ad-cb > b - d >= 0
if (t == 1 && other.t == 1) return p < other.p;
if (t == 1) return false;
else if (other.t == 1) return true;
return (1-other.t) * (t * -p) < (1-t) * (other.t * -other.p);
// return t == other.t ? p < other.p : t < other.t;
}
};
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
ll a = A, n = P.size();
ll required = 0, one = 0, two = 0, three = 0, four = 0;
vector<coupon> coupons[5];
vector<coupon> cps;
for (int i = 0; i < n; i++) {
coupons[T[i]].push_back(coupon{T[i], P[i], i}), required += P[i];
cps.push_back(coupons[T[i]].back());
if (T[i] == 1) one++;
else if (T[i] == 2) two++;
else if (T[i] == 3) three++;
else four++;
}
sort(all(cps));
// cout << one << " " << two << " " << three << " " << four << "\n";
for (int i = 1; i <= 4; i++) sort(all(coupons[i]));
vector<ll> pref1; for (int i = 0; i < coupons[1].size(); i++) {
if (i == 0) pref1.push_back(coupons[1][i].p);
else pref1.push_back(pref1.back() + coupons[1][i].p);
}
// for (auto &el : cps) cout << el.idx << " "; cout << "\n";
vector<int> best, pos;
int b1 = 0, c1 = upper_bound(all(pref1), a) - pref1.begin();
vector<int> unused;
for (int i = 0; i < n; i++) {
if (cps[i].t == 1) break;
if (cps[i].norm(a)) {
unused.push_back(i);
continue;
}
if (cps[i].eval(a)<0) break;
a = cps[i].eval(a);
best.push_back(cps[i].idx);
}
int m = unused.size(); int k = min(m, 20);
vector<ll> dp(k+1, -4e18); vector<vector<ll>> s(k+1); dp[0] = a;
for (int i = 0; i < m; i++) {
for (int j = k-1; j >= 0; j--) {
ll f = cps[unused[i]].eval(dp[j]);
if (f < 0) continue;
if (f < dp[j+1]) continue;
s[j+1] = s[j]; s[j+1].push_back(unused[i]);
dp[j+1] = f;
}
}
int b = 0, c = upper_bound(all(pref1), a) - pref1.begin();
for (int i = 1; i <= k; i++) {
int c1 = upper_bound(all(pref1), dp[i]) - pref1.begin();
if (s[i].size() + c1 > s[b].size() + c) {
b = i; c = c1;
}
}
for (auto &el : s[b]) best.push_back(cps[el].idx);
for (int i = 0; i < c; i++) best.push_back(coupons[1][i].idx);
// int x = 0, y = 0, z = 0;
// for (int i = 1; i <= two+three+four; i++) {
// coupon b; bool found = false; int which = -1;
// if (x<two && !found || coupons[2][x] < b) found = true, b = coupons[2][x], which = 2;
// if (y<three && !found || coupons[3][y] < b) found = true, b = coupons[3][y], which = 3;
// if (z<three && !found || coupons[4][z] < b) found = true, b = coupons[4][z], which = 4;
// cout << "Chose " << b.idx << "\n";
// a = b.eval(a);
// if (a < 0) break;
// if (which == 2) pos.push_back(coupons[which][x++].idx);
// else if (which == 3) pos.push_back(coupons[which][y++].idx);
// else if (which == 4) pos.push_back(coupons[which][z++].idx);
// int c2 = upper_bound(all(pref1), a) - pref1.begin();
// if (pos.size() + c2 > best.size() + c1) {
// best = pos; c1 = c2;
// }
// }
// for (int i = 0; i < c1; i++) best.push_back(coupons[1][i].idx);
// vector<vector<vector<ll>>> DP(two+1, vector<vector<ll>>(three+1, vector<ll>(four+1, -4e18))), par(two+1, vector<vector<ll>>(three+1, vector<ll>(four+1, -1)));
// DP[0][0][0] = A;
// for (int x = 0; x <= two; x++) {
// for (int y = 0; y <= three; y++) {
// for (int z = 0; z <= four; z++) {
// if (DP[x][y][z] < 0) continue;
// if (x<two && DP[x+1][y][z] < coupons[2][x].eval(DP[x][y][z])) {
// DP[x+1][y][z] = coupons[2][x].eval(DP[x][y][z]);
// par[x+1][y][z] = 2;
// }
// if (y < three && DP[x][y+1][z] < coupons[3][y].eval(DP[x][y][z])) {
// DP[x][y+1][z] = coupons[3][y].eval(DP[x][y][z]);
// par[x][y+1][z] = 3;
// }
// if (z < four && DP[x][y][z+1] < coupons[4][z].eval(DP[x][y][z])) {
// DP[x][y][z+1] = coupons[4][z].eval(DP[x][y][z]);
// par[x][y][z+1] = 4;
// }
// }
// }
// }
// for (int i = 0; i <= two; i++) {
// for (int j = 0; j <= three; j++) {
// for (int k = 0; k <= four; k++) {
// if (DP[i][j][k] < 0) continue;
// int x = i, y = j, z = k;
// vector<int> trail;
// while (par[x][y][z] != -1) {
// trail.push_back(par[x][y][z]);
// if (par[x][y][z] == 2) x--;
// else if (par[x][y][z] == 3) y--;
// else if (par[x][y][z] == 4) z--;
// }
// reverse(trail.begin(), trail.end());
// x = y = z = 0;
// vector<int> ans;
// for (int u = 0; u < trail.size(); u++) {
// if (trail[u] == 2) ans.push_back(coupons[trail[u]][x++].idx);
// else if (trail[u] == 3) ans.push_back(coupons[trail[u]][y++].idx);
// else if (trail[u] == 4) ans.push_back(coupons[trail[u]][z++].idx);
// }
// int c1 = upper_bound(all(pref1), DP[i][j][k]) - pref1.begin();
// for (int u = 0; u < c1; u++) ans.push_back(coupons[1][u].idx);
// if (ans.size() >= best.size()) {
// swap(ans, best);
// }
// }
// }
// }
return best;
}
