#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
struct dp_helper {
};
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int n = P.size();
vector<vector<pair<int, int>>> can(4);
for (int i = 0; i < n; i++) {
T[i]--;
can[T[i]].emplace_back(P[i], i);
}
for (int i = 0; i < 4; i++) {
ranges::sort(can[i]);
}
vector<ll> pref(can[0].size() + 1, 0);
for (int i = 0; i < can[0].size(); i++) {
pref[i + 1] = pref[i] + can[0][i].first;
}
map<tuple<int, int, int>, pair<ll, int>> dp_cur, dp_next;
map<tuple<int, int, int>, pair<ll, int>> saved;
dp_cur[{0, 0, 0}] = {A, -1};
auto update_with_incr = [&](int i, int j, int l, int tr) -> void {
ll res = dp_cur[{i, j, l}].first;
if (tr == 1) {
if (can[1].size() > i && (res - can[1][i].first) * 2 > res) {
dp_next[{i + 1, j, l}] = max(dp_next[{i + 1, j, l}], {(res - can[1][i].first) * 2, 1});
}
} else if (tr == 2) {
if (can[2].size() > j && (res - can[2][j].first) * 3 > res) {
dp_next[{i, j + 1, l}] = max(dp_next[{i, j + 1, l}], {(res - can[2][j].first) * 3, 2});
}
} else {
if (can[3].size() > l && (res - can[3][l].first) * 4 > res) {
dp_next[{i, j, l + 1}] = max(dp_next[{i, j, l + 1}], {(res - can[3][l].first) * 4, 3});
}
}
};
auto update_with = [&](int i, int j, int l, int tr) -> void {
ll res = dp_cur[{i, j, l}].first;
if (tr == 1) {
if (can[1].size() > i && res - can[1][i].first >= 0) {
dp_next[{i + 1, j, l}] = max(dp_next[{i + 1, j, l}], {(res - can[1][i].first) * 2, 1});
}
} else if (tr == 2) {
if (can[2].size() > j && res - can[2][j].first >= 0) {
dp_next[{i, j + 1, l}] = max(dp_next[{i, j + 1, l}], {(res - can[2][j].first) * 3, 2});
}
} else {
if (can[3].size() > l && res - can[3][l].first >= 0) {
dp_next[{i, j, l + 1}] = max(dp_next[{i, j, l + 1}], {(res - can[3][l].first) * 4, 3});
}
}
};
auto max_ans = [&](int i, int j, int l) -> int {
int l1 = 0, r1 = can[0].size();
ll have = dp_cur[{i, j, l}].first;
while (l1 != r1) {
int m = (l1 + r1 + 1) / 2;
if (pref[m] <= have) {
l1 = m;
} else {
r1 = m - 1;
}
}
return i + j + l + l1;
};
int ans = max_ans(0, 0, 0);
tuple<int, int, int> ans_coords = {0, 0, 0};
for (int now_sum = 0; now_sum < 32; now_sum++) {
for (auto [k, v] : dp_cur) {
for (int tr = 1; tr < 4; tr++) {
update_with(get<0>(k), get<1>(k), get<2>(k), tr);
}
saved[k] = v;
}
swap(dp_cur, dp_next);
dp_next.clear();
for (auto [k, v] : dp_cur) {
int ma = max_ans(get<0>(k), get<1>(k), get<2>(k));
if (ma > ans) {
ans = ma;
ans_coords = k;
}
}
}
vector<int> vans;
ll have = saved[ans_coords].first;
while (ans_coords != make_tuple(0, 0, 0)) {
int tr = saved[ans_coords].second;
if (tr == 1) {
vans.push_back(can[1][get<0>(ans_coords) - 1].second);
get<0>(ans_coords)--;
} else if (tr == 2) {
vans.push_back(can[2][get<1>(ans_coords) - 1].second);
get<1>(ans_coords)--;
} else {
vans.push_back(can[3][get<2>(ans_coords) - 1].second);
get<2>(ans_coords)--;
}
}
ranges::reverse(vans);
int ptr = 0;
while (ptr < can[0].size() && have >= can[0][ptr].first) {
have -= can[0][ptr].first;
vans.push_back(can[0][ptr].second);
ptr++;
}
return vans;
}
