#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define all(x) x.begin(), x.end()
constexpr ll INF = ll(1e9) * 200010ll;
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
const int n = P.size();
vector<vector<pair<ll, int>>> L(4);
ll sum = 0;
for (int i = 0; i < n; i++) {
T[i]--;
L[T[i]].push_back({P[i], i});
sum += P[i];
}
for (int i = 0; i < 4; i++) {
sort(all(L[i]));
}
vector<vector<vector<vector<pair<ll, int>>>>> dp(L[0].size() + 1, // Max remaining coupons, last
vector<vector<vector<pair<ll, int>>>>(L[1].size() + 1,
vector<vector<pair<ll, int>>>(L[2].size() + 1,
vector<pair<ll, int>>(L[3].size() + 1, {-1, 0}))));
array<int, 4> best = {0, 0, 0, 0};
int mx = 0;
dp[0][0][0][0] = {A, -1};
for (int i = 0; i <= int(L[0].size()); i++) {
for (int j = 0; j <= int(L[1].size()); j++) {
for (int k = 0; k <= int(L[2].size()); k++) {
for (int h = 0; h <= int(L[3].size()); h++) {
auto [x, _] = dp[i][j][k][h];
if (x == -1) continue;
if (i + j + k + h > mx) {
mx = i + j + k + h;
best = {i, j, k, h};
}
if (i + 1 <= int(L[0].size()) && x >= L[0][i].first) {
dp[i + 1][j][k][h] = max(dp[i + 1][j][k][h], {min(INF, (x - L[0][i].first) * 1ll), 0});
}
if (j + 1 <= int(L[1].size()) && x >= L[1][j].first) {
dp[i][j + 1][k][h] = max(dp[i][j + 1][k][h], {min(INF, (x - L[1][j].first) * 2ll), 1});
}
if (k + 1 <= int(L[2].size()) && x >= L[2][k].first) {
dp[i][j][k + 1][h] = max(dp[i][j][k + 1][h], {min(INF, (x - L[2][k].first) * 3ll), 2});
}
if (h + 1 <= int(L[3].size()) && x >= L[3][h].first) {
dp[i][j][k][h + 1] = max(dp[i][j][k][h + 1], {min(INF, (x - L[3][h].first) * 4ll), 3});
}
}
}
}
}
vector<int> res;
while (true) {
auto &[i, j, k, h] = best;
auto [x, last] = dp[i][j][k][h];
if (last == -1) {
break;
}
if (last == 0) {
res.push_back(L[0][--i].second);
} else if (last == 1) {
res.push_back(L[1][--j].second);
} else if (last == 2) {
res.push_back(L[2][--k].second);
} else if (last == 3) {
res.push_back(L[3][--h].second);
}
}
reverse(all(res));
return res;
}
//#include "grader.cpp"
