#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
std::vector<int> max_coupons(int _A, std::vector<int> P, std::vector<int> T) {
int N = P.size();
vector<vector<int>> coupons(5);
for(int i=0; i<N; i++) {
coupons[T[i]].push_back(i);
}
for(int i=0; i<5; i++) sort(coupons[i].begin(), coupons[i].end(), [&](int a, int b) {
return P[a] < P[b];
});
int A = coupons[2].size();
int B = coupons[3].size();
int C = coupons[4].size();
long long money = _A;
long long sum = accumulate(P.begin(), P.end(), 0ll);
vector dp(A+1, vector(B+1, vector<array<long long,4>>(C+1, {-1, -1, -1, -1})));
dp[0][0][0] = {_A, -1, -1, -1};
for(int a=0; a<=A; a++) {
for(int b=0; b<=B; b++) {
for(int c=0; c<=C; c++) {
long long x = dp[a][b][c][0];
// cerr << a << " " << b << " " << c << " " << x << endl;
if(dp[a][b][c][0] == -1) continue;
if(a < A and x >= P[coupons[2][a]]) {
long long new_money = min(sum, (x - P[coupons[2][a]]) * 2);
dp[a+1][b][c] = max(dp[a+1][b][c], array<long long, 4>({new_money, a, b, c}));
}
if(b < B and x >= P[coupons[3][b]]) {
long long new_money = min(sum, (x - P[coupons[3][b]]) * 3);
dp[a][b+1][c] = max(dp[a][b+1][c], array<long long, 4>({new_money, a, b, c}));
}
if(c < C and x >= P[coupons[4][c]]) {
long long new_money = min(sum, (x - P[coupons[4][c]]) * 4);
dp[a][b][c+1] = max(dp[a][b][c+1], array<long long, 4>({new_money, a, b, c}));
}
}
}
}
vector<long long> pfs = {0};
for(int x: coupons[1]) {
pfs.push_back(pfs.back() + P[x]);
}
long long score = 0;
array<int, 3> state = {0, 0, 0};
for(int a=0; a<=A; a++) {
for(int b=0; b<=B; b++) {
for(int c=0; c<=C; c++) {
if(dp[a][b][c][0] == -1) continue;
long long bought = a + b + c;
long long money = dp[a][b][c][0];
for(int x: coupons[1]) {
if(money < P[x]) break;
money -= P[x];
bought++;
}
if(bought > score) {
score = bought;
state = {a, b, c};
}
}
}
}
vector<int> bought;
auto [a, b, c] = state;
while(a > 0 or b > 0 or c > 0) {
auto [x, na, nb, nc] = dp[a][b][c];
int curr = -1;
if(na != a) curr = coupons[2][na];
if(nb != b) curr = coupons[3][nb];
if(nc != c) curr = coupons[4][nc];
a = na;
b = nb;
c = nc;
bought.push_back(curr);
}
reverse(bought.begin(), bought.end());
money = dp[state[0]][state[1]][state[2]][0];
for(int x: coupons[1]) {
if(money < P[x]) break;
money -= P[x];
bought.push_back(x);
}
return bought;
}
