#include "festival.h"
#include <algorithm>
#include <iostream>
#include <numeric>
#include <vector>
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
std::vector<std::vector<int>> coup(4);
int N = P.size();
for (int i = 0; i < N; ++i) {
coup[T[i] - 1].push_back(i);
}
for (int i = 0; i < 4; ++i) {
std::sort(coup[i].begin(), coup[i].end(),
[&](int x, int y) { return P[x] > P[y]; });
}
std::vector<int> ans;
int64_t total = std::accumulate(P.begin(), P.end(), 0LL);
int64_t X = A;
while (true) {
int j = -1;
int64_t Z = X;
// std::cerr << "X = " << X << "\n";
for (int i = 0; i < 4; ++i) {
if (coup[i].empty() || X < P[coup[i].back()]) {
continue;
}
int64_t Y = (X - P[coup[i].back()]) * (i + 1);
// std::cerr << "type " << i + 1 << " Y = " << Y << "\n";
if (Y > Z) {
Z = Y;
j = i;
}
}
if (j == -1) {
break;
}
X = std::min((X - P[coup[j].back()]) * (j + 1), total);
ans.push_back(coup[j].back());
coup[j].pop_back();
}
std::vector<int> rem;
for (int i = 0; i < 4; ++i) {
rem.insert(rem.end(), coup[i].begin(), coup[i].end());
}
sort(rem.begin(), rem.end(),
[&](int i, int j) { return P[i] * T[j] < T[i] * P[j]; });
for (int i : rem) {
if (X < P[i]) continue;
X -= P[i];
X *= T[i];
ans.push_back(i);
}
return ans;
}
