#include "festival.h"
#include <algorithm>
#include <vector>
#include <cstdint>
using namespace std;
using ll = long long;
using i128 = __int128_t;
const ll INF = 4'000'000'000'000'000'000LL; // 4e18
struct Booster {
int idx; // original index
int P;
int T;
};
struct T1 {
ll P;
int idx;
};
struct State {
int cnt; // number of boosters taken
ll X; // current tokens (capped at INF)
int prev; // index of previous state in the global vector
int taken; // index in the boosters array of the booster taken to reach this state, or -1
};
// comparator for sorting boosters by key = P * T / (T-1)
bool booster_cmp(const Booster &a, const Booster &b) {
// compare a.P * a.T * (b.T - 1) vs b.P * b.T * (a.T - 1)
i128 left = (i128)a.P * a.T * (b.T - 1);
i128 right = (i128)b.P * b.T * (a.T - 1);
if (left != right) return left < right;
// tie‑break by original index (any order works)
return a.idx < b.idx;
}
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int N = P.size();
// split coupons
vector<Booster> boosters;
vector<T1> t1;
for (int i = 0; i < N; ++i) {
if (T[i] > 1) {
boosters.push_back({i, P[i], T[i]});
} else {
t1.push_back({P[i], i});
}
}
// sort boosters by the optimal key
sort(boosters.begin(), boosters.end(), booster_cmp);
// sort T=1 coupons by cost
sort(t1.begin(), t1.end(),
[](const T1 &a, const T1 &b) { return a.P < b.P; });
// prefix sums of T1 costs
vector<ll> pref_t1;
for (auto &x : t1) pref_t1.push_back(x.P);
for (size_t i = 1; i < pref_t1.size(); ++i)
pref_t1[i] += pref_t1[i - 1];
// DP: we store all states that ever appear.
// The states are kept in a global vector; only the frontier (Pareto‑optimal)
// is kept for the current prefix.
vector<State> states;
states.push_back({0, (ll)A, -1, -1}); // initial state (no booster taken)
vector<int> cur_idx = {0}; // indices of states in the current frontier
// process boosters in the sorted order
for (size_t bi = 0; bi < boosters.size(); ++bi) {
const Booster &b = boosters[bi];
vector<int> new_idxs;
// generate states obtained by taking the current booster
for (int idx : cur_idx) {
const State &s = states[idx];
if (s.X >= b.P) {
i128 newX_i128 = (i128)(s.X - b.P) * b.T;
ll newX = (newX_i128 > INF) ? INF : (ll)newX_i128;
int new_idx = (int)states.size();
states.push_back({s.cnt + 1, newX, idx, (int)bi});
new_idxs.push_back(new_idx);
}
}
// Merge the two sets (skip = cur_idx, take = new_idxs) into a new frontier.
// Both lists are already sorted by cnt increasing and X strictly decreasing.
vector<tuple<int, ll, int>> best; // (cnt, X, idx) with best X for each cnt
size_t i = 0, j = 0;
while (i < cur_idx.size() || j < new_idxs.size()) {
int cnt;
ll X;
int idx;
if (j == new_idxs.size() ||
(i < cur_idx.size() && states[cur_idx[i]].cnt < states[new_idxs[j]].cnt)) {
cnt = states[cur_idx[i]].cnt;
X = states[cur_idx[i]].X;
idx = cur_idx[i];
++i;
}
else if (i == cur_idx.size() ||
(j < new_idxs.size() && states[cur_idx[i]].cnt > states[new_idxs[j]].cnt)) {
cnt = states[new_idxs[j]].cnt;
X = states[new_idxs[j]].X;
idx = new_idxs[j];
++j;
}
else {
// same cnt – keep the one with larger X
const State &s1 = states[cur_idx[i]];
const State &s2 = states[new_idxs[j]];
if (s1.X >= s2.X) {
cnt = s1.cnt;
X = s1.X;
idx = cur_idx[i];
} else {
cnt = s2.cnt;
X = s2.X;
idx = new_idxs[j];
}
++i; ++j;
}
best.emplace_back(cnt, X, idx);
}
// Remove dominated states: keep only those with strictly decreasing X as cnt increases.
vector<int> next_idx;
for (auto &[cnt, X, idx] : best) {
while (!next_idx.empty() && X >= states[next_idx.back()].X) {
next_idx.pop_back();
}
if (next_idx.empty() || X < states[next_idx.back()].X) {
next_idx.push_back(idx);
}
}
cur_idx = move(next_idx);
}
// Evaluate all final states to choose the best total number of coupons.
int best_state = -1;
ll best_total = -1;
for (int idx : cur_idx) {
const State &s = states[idx];
int t = 0;
if (!t1.empty()) {
t = upper_bound(pref_t1.begin(), pref_t1.end(), s.X) - pref_t1.begin();
}
ll total = s.cnt + t;
if (total > best_total) {
best_total = total;
best_state = idx;
}
}
// Reconstruct the sequence of purchases.
// First, the boosters in chronological order.
vector<int> result;
int cur = best_state;
while (cur != -1) {
if (states[cur].taken != -1) {
// taken a booster: record its original index
result.push_back(boosters[states[cur].taken].idx);
}
cur = states[cur].prev;
}
reverse(result.begin(), result.end());
// Then the T=1 coupons (cheapest first).
if (!t1.empty()) {
ll X_final = states[best_state].X;
int t = upper_bound(pref_t1.begin(), pref_t1.end(), X_final) - pref_t1.begin();
for (int i = 0; i < t; ++i) {
result.push_back(t1[i].idx);
}
}
return result;
}
