#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef __int128 i128;
const ll INF = 2'000'000'000'000'000'000LL; // 2e18, safely above any relevant sum
struct Coupon {
int idx;
int P;
int T;
ll a, d; // a = T, d = T*P
i128 num, den; // for sorting: num = P * T, den = T-1 (as i128)
};
struct Type1 {
int idx;
ll P;
};
struct Node {
int m; // number of taken type>1 coupons
ll B; // current budget = M * (A - S0)
int prev; // index of previous node in nodes vector
int last_idx; // index of the coupon (original) that led to this node, -1 for none
};
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
int N = P.size();
vector<Coupon> type2;
vector<Type1> type1;
for (int i = 0; i < N; ++i) {
if (T[i] == 1) {
type1.push_back({i, P[i]});
} else {
type2.push_back({i, P[i], T[i], (ll)T[i], (ll)T[i] * P[i]});
type2.back().num = (i128)P[i] * T[i];
type2.back().den = T[i] - 1;
}
}
// sort type2 by key = (P*T)/(T-1) increasing
sort(type2.begin(), type2.end(), [](const Coupon& x, const Coupon& y) {
i128 left = x.num * y.den;
i128 right = y.num * x.den;
if (left != right) return left < right;
// break ties by T descending (higher multiplier first) – not critical
return x.T > y.T;
});
// sort type1 by price ascending (optimal for buying after all type2)
sort(type1.begin(), type1.end(), [](const Type1& a, const Type1& b) {
return a.P < b.P;
});
// prefix sums of type1 prices
vector<ll> pref(type1.size() + 1, 0);
for (size_t i = 0; i < type1.size(); ++i)
pref[i+1] = pref[i] + type1[i].P;
// DP over type2 coupons using Pareto frontier
vector<Node> nodes;
nodes.push_back({0, A, -1, -1}); // initial state (taken 0, budget A)
vector<int> frontier = {0}; // indices of Pareto-optimal states
for (const Coupon& cp : type2) {
vector<int> new_candidates;
for (int idx : frontier) {
const Node& cur = nodes[idx];
if (cur.B >= cp.P) {
ll newB;
if (cur.B > INF / cp.a) {
newB = INF; // cap too large values
} else {
i128 tmp = (i128)cp.a * cur.B - cp.d;
newB = (tmp > INF) ? INF : (ll)tmp;
}
nodes.push_back({cur.m + 1, newB, idx, cp.idx});
new_candidates.push_back(nodes.size() - 1);
}
}
// combine old frontier and new candidates, then keep only Pareto-optimal states
vector<int> all = frontier;
all.insert(all.end(), new_candidates.begin(), new_candidates.end());
// sort by m ascending, then by B descending (so that for equal m we keep the highest B)
sort(all.begin(), all.end(), [&](int i, int j) {
if (nodes[i].m != nodes[j].m) return nodes[i].m < nodes[j].m;
return nodes[i].B > nodes[j].B;
});
// keep only one state per m (the one with highest B)
vector<int> unique_m;
for (int i : all) {
if (unique_m.empty() || nodes[i].m != nodes[unique_m.back()].m)
unique_m.push_back(i);
}
// now keep only points with strictly decreasing B as m increases
vector<int> new_frontier;
for (int i : unique_m) {
while (!new_frontier.empty() && nodes[i].B >= nodes[new_frontier.back()].B)
new_frontier.pop_back();
new_frontier.push_back(i);
}
frontier = move(new_frontier);
}
// find the best state (max total coupons)
ll best_total = -1;
int best_node = -1;
for (int idx : frontier) {
const Node& nd = nodes[idx];
int t = upper_bound(pref.begin(), pref.end(), nd.B) - pref.begin() - 1;
ll total = nd.m + t;
if (total > best_total) {
best_total = total;
best_node = idx;
}
}
// reconstruction: first build list of taken type2 in reverse order
vector<int> answer;
int cur = best_node;
while (cur != -1) {
if (nodes[cur].last_idx != -1)
answer.push_back(nodes[cur].last_idx);
cur = nodes[cur].prev;
}
// answer now contains taken type2 in reverse chronological order; reverse it
reverse(answer.begin(), answer.end());
// add the smallest type1 coupons
const Node& best = nodes[best_node];
int t = upper_bound(pref.begin(), pref.end(), best.B) - pref.begin() - 1;
for (int i = 0; i < t; ++i)
answer.push_back(type1[i].idx);
return answer;
}
