#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int N = (int)P.size();
// Separate coupons by type, keep original indices
vector<pair<long long,int>> byType[5];
for (int i = 0; i < N; ++i) {
byType[T[i]].push_back({ (long long)P[i], i });
}
for (int t = 1; t <= 4; ++t)
sort(byType[t].begin(), byType[t].end(),
[](const pair<long long,int>& a, const pair<long long,int>& b){
return a.first < b.first;
});
// Extract price and index vectors for each type (ascending prices)
vector<long long> price1, idx1;
vector<long long> price2, idx2;
vector<long long> price3, idx3;
vector<long long> price4, idx4;
for (auto &pr : byType[1]) { price1.push_back(pr.first); idx1.push_back(pr.second); }
for (auto &pr : byType[2]) { price2.push_back(pr.first); idx2.push_back(pr.second); }
for (auto &pr : byType[3]) { price3.push_back(pr.first); idx3.push_back(pr.second); }
for (auto &pr : byType[4]) { price4.push_back(pr.first); idx4.push_back(pr.second); }
int n1 = (int)price1.size();
int n2 = (int)price2.size();
int n3 = (int)price3.size();
int n4 = (int)price4.size();
// DP over counts of types 2,3,4
const long long INF = (1LL<<60); // large enough (≈1e18)
int dim3 = n3 + 1;
int dim4 = n4 + 1;
size_t totalStates = (size_t)(n2 + 1) * dim3 * dim4;
vector<long long> dp(totalStates, -1); // -1 means unreachable
vector<char> pre(totalStates, -1); // which type was added to reach the state
auto index = [&](int i2, int i3, int i4) -> size_t {
return ((size_t)i2 * dim3 + i3) * dim4 + i4;
};
dp[index(0,0,0)] = A; // start with all tokens
for (int i2 = 0; i2 <= n2; ++i2) {
for (int i3 = 0; i3 <= n3; ++i3) {
for (int i4 = 0; i4 <= n4; ++i4) {
size_t curIdx = index(i2,i3,i4);
long long cur = dp[curIdx];
if (cur < 0) continue;
// take a type‑2 coupon
if (i2 < n2) {
long long price = price2[i2];
if (cur >= price) {
__int128 after = (__int128)(cur - price) * 2;
long long nxt = after > INF ? INF : (long long)after;
size_t nxtIdx = index(i2+1,i3,i4);
if (nxt > dp[nxtIdx]) {
dp[nxtIdx] = nxt;
pre[nxtIdx] = 2;
}
}
}
// take a type‑3 coupon
if (i3 < n3) {
long long price = price3[i3];
if (cur >= price) {
__int128 after = (__int128)(cur - price) * 3;
long long nxt = after > INF ? INF : (long long)after;
size_t nxtIdx = index(i2,i3+1,i4);
if (nxt > dp[nxtIdx]) {
dp[nxtIdx] = nxt;
pre[nxtIdx] = 3;
}
}
}
// take a type‑4 coupon
if (i4 < n4) {
long long price = price4[i4];
if (cur >= price) {
__int128 after = (__int128)(cur - price) * 4;
long long nxt = after > INF ? INF : (long long)after;
size_t nxtIdx = index(i2,i3,i4+1);
if (nxt > dp[nxtIdx]) {
dp[nxtIdx] = nxt;
pre[nxtIdx] = 4;
}
}
}
}
}
}
// Prefix sums for type‑1 coupons (sorted increasingly)
vector<long long> pref1(n1 + 1, 0);
for (int i = 0; i < n1; ++i) pref1[i+1] = pref1[i] + price1[i];
// Find best reachable state (max total coupons)
int bestTotal = -1;
int best_i2 = 0, best_i3 = 0, best_i4 = 0;
int best_k1 = 0;
for (int i2 = 0; i2 <= n2; ++i2) {
for (int i3 = 0; i3 <= n3; ++i3) {
for (int i4 = 0; i4 <= n4; ++i4) {
long long cur = dp[index(i2,i3,i4)];
if (cur < 0) continue;
// maximum number of type‑1 coupons we can still afford
int k1 = (int)(upper_bound(pref1.begin(), pref1.end(), cur) - pref1.begin()) - 1;
int total = i2 + i3 + i4 + k1;
if (total > bestTotal) {
bestTotal = total;
best_i2 = i2; best_i3 = i3; best_i4 = i4;
best_k1 = k1;
}
}
}
}
// Reconstruct the ordering of types 2‑4 coupons
vector<int> answer;
int c2 = best_i2, c3 = best_i3, c4 = best_i4;
while (c2 > 0 || c3 > 0 || c4 > 0) {
size_t curIdx = index(c2,c3,c4);
char t = pre[curIdx];
if (t == -1) break; // safety, should not happen
if (t == 2) {
answer.push_back((int)idx2[c2-1]);
--c2;
} else if (t == 3) {
answer.push_back((int)idx3[c3-1]);
--c3;
} else { // t == 4
answer.push_back((int)idx4[c4-1]);
--c4;
}
}
// The vector now holds coupons in reverse order (last bought first)
reverse(answer.begin(), answer.end());
// Append the cheapest type‑1 coupons we can afford
for (int i = 0; i < best_k1; ++i)
answer.push_back((int)idx1[i]);
return answer;
}
