#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
#define int long long
int INF = 1LL<<60;
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;
};
//first 3 subtasks
int best_overall_count = -1;
int best_type_2_count = -1;
int money = A;
for (int i=0; i<=coupons[2].size(); i++) {
int type_2_count = i;
int type_1_count = get_num_type_1(money);
int totcount = type_2_count + type_1_count;
if (totcount > best_overall_count) {
best_overall_count = totcount;
best_type_2_count = type_2_count;
}
if (i<coupons[2].size()) {
money = (money - coupons[2][i].first) * 2;
}
if (money <= 0) break;
money = min(money, INF);
}
//cout << "SOLVED " << best_type_2_count << " " << best_overall_count << endl;
vector<signed> answers;
for (int i=0; i<best_type_2_count; i++) {
answers.push_back(coupons[2][i].second);
}
for (int i=0; i<best_overall_count-best_type_2_count; i++) {
answers.push_back(coupons[1][i].second);
}
return answers;
}
