#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
#define int long long
int INF = 1LL<<60;
int M = 75;
vector<signed> answers;
std::vector<signed> max_coupons(signed A, std::vector<signed> P, std::vector<signed> T) {
//cost, index
int n = P.size();
vector<vector<pair<int,int>>> coupons(5);
for (int i=0; i<n; i++) {
coupons[T[i]].push_back({P[i], i});
}
for (int i=1; i<5; i++) {
sort(coupons[i].begin(), coupons[i].end());
}
vector<vector<int>> prefsums(5);
for (int i=1; i<5; i++) {
prefsums[i] = vector<int>(coupons[i].size() + 1, 0);
for (int j=0; j<coupons[i].size(); j++) {
prefsums[i][j+1] = prefsums[i][j] + coupons[i][j].first;
}
}
auto get_num_type_1 = [&prefsums] (int money) {
return upper_bound(prefsums[1].begin(), prefsums[1].end(), money) - prefsums[1].begin() - 1;
};
//dp[i][j][k] = max money reachable from using the first i 2-coupons, first j 3-coupons and first k 4-coupons
int n2 = coupons[2].size();
int n3 = coupons[3].size();
int n4 = coupons[4].size();
int dp[n2+1][n3+1][n4+1];
int back[n2+1][n3+1][n4+1];
for (int i=0; i<=n2; i++) {
for (int j=0; j<=n3; j++) {
for (int k=0; k<=n4; k++) {
dp[i][j][k] = -INF;
back[i][j][k] = -1; //x if we just used a coupon x
}
}
}
dp[0][0][0] = A;
for (int i=0; i<=n2; i++) {
for (int j=0; j<=n3; j++) {
for (int k=0; k<=n4; k++) {
int cell = dp[i][j][k];
if (cell <= 0) {
continue;
}
if (i<n2) {
if (((cell - coupons[2][i].first) * 2) > dp[i+1][j][k]) {
dp[i+1][j][k] = min(INF, ((cell - coupons[2][i].first) * 2));
back[i+1][j][k] = 2;
}
}
if (j<n3) {
if (((cell - coupons[3][j].first) * 3) > dp[i][j+1][k]) {
dp[i][j+1][k]= min(INF, ((cell - coupons[3][j].first) * 3));
back[i][j+1][k] = 3;
}
}
if (k<n4) {
if (((cell - coupons[4][k].first) * 4) > dp[i][j][k+1]) {
dp[i][j][k+1]= min(INF, ((cell - coupons[4][k].first) * 4));
back[i][j][k+1] = 4;
}
}
}
}
}
int maxres = -1;
vector<int> maxpos;
for (int i=0; i<=n2; i++) {
for (int j=0; j<=n3; j++) {
for (int k=0; k<=n4; k++) {
if (dp[i][j][k]<0) continue;
int res = i + j + k + get_num_type_1(dp[i][j][k]);
if (res > maxres) {
maxres = res;
maxpos = {i, j, k};
}
}
}
}
vector<int> curpos(maxpos.begin(), maxpos.end());
vector<int> end = {0,0,0};
while (curpos != end) {
int pter = back[curpos[0]][curpos[1]][curpos[2]];
//cout << curpos[0] <<" " << curpos[1] <<" " << curpos[2] <<" " << pter << endl;
curpos[pter-2]--;
answers.push_back(coupons[pter][curpos[pter-2]].second);
}
reverse(answers.begin(), answers.end());
for (int i=0; i<get_num_type_1(dp[maxpos[0]][maxpos[1]][maxpos[2]]); i++) {
answers.push_back(coupons[1][i].second);
}
return answers;
}
