#include "festival.h"
#include <bits/stdc++.h>
#define cost first
#define ind second
using namespace std;
int c = 0;
vector<int> pref1;
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
vector<vector<pair<int, int>>> options(4, vector<pair<int, int>>{});
for (int i = 0; i < P.size(); ++i) {
options[T[i]-1].push_back({P[i], i});
}
for (int i = 0; i < 4; ++i) sort(options[i].begin(), options[i].end());
for (int i = 0; i < options[0].size(); ++i) {
for (int j = 0; j < options[0][i].first; ++j) pref1.push_back(c);
++c;
}
int t2 = options[1].size(), t3 = options[2].size(), t4 = options[3].size();
pair<int, vector<int>> dp[t2+1][t3+1][t4+1];
for (int i = 0; i < t2; ++i) {
for (int j = 0; j < t3; ++j) {
for (int k = 0; k < t4; ++k) {
dp[i][j][k] = {-1, vector<int>{}};
}
}
}
dp[0][0][0] = {A, vector<int>{}};
for (int i = 0; i <= t2; ++i) {
for (int j = 0; j <= t3; ++j) {
for (int k = 0; k <= t4; ++k) {
if (dp[i][j][k].first == -1) continue;
if (i != t2) {
if ((dp[i][j][k].first - options[1][i].cost) * 2 > dp[i+1][j][k].first && options[1][i].cost <= dp[i][j][k].first) {
dp[i+1][j][k] = {(dp[i][j][k].first - options[1][i].cost) * 2, dp[i][j][k].second};
dp[i+1][j][k].second.push_back(options[1][i].ind);
}
}
if (j != t3) {
if ((dp[i][j][k].first - options[2][j].cost) * 3 > dp[i][j+1][k].first && options[2][j].cost <= dp[i][j][k].first) {
dp[i][j+1][k] = {(dp[i][j][k].first - options[2][j].cost) * 3, dp[i][j][k].second};
dp[i][j+1][k].second.push_back(options[2][j].ind);
}
}
if (k != t4) {
if ((dp[i][j][k].first - options[3][k].cost) * 4 > dp[i][j][k+1].first && options[3][k].cost <= dp[i][j][k].first) {
dp[i][j][k+1] = {(dp[i][j][k].first - options[3][k].cost) * 4, dp[i][j][k].second};
dp[i][j][k+1].second.push_back(options[3][k].ind);
}
}
cout<<i<<' '<<j<<' '<<k<<' '<<dp[i][j][k].first<<'\n';
}
}
}
int a, b, o, d;
vector<int> ret;
for (int i = 0; i <= t2; ++i) {
for (int j = 0; j <= t3; ++j) {
for (int k = 0; k <= t4; ++k) {
if (dp[i][j][k].first == -1) continue;
int n = i + j + k;
if (dp[i][j][k].first >= pref1.size()) {
ret = dp[i][j][k].second;
for (int l = 0; l < c; ++l) {
ret.push_back(options[0][l].ind);
}
}
int _c = pref1[dp[i][j][k].first];
if (ret.size() < n + _c) {
ret = dp[i][j][k].second;
for (int l = 0; l < _c; ++l) {
ret.push_back(options[0][l].ind);
}
}
}
}
}
return ret;
}
