#include "festival.h" #include <bits/stdc++.h> using namespace std; #define all(x) (x).begin(), (x).end() typedef long long ll; using pi = pair<ll, ll>; using vi = vector<ll>; struct coupon { ll t, p, idx; ll eval(ll R) { if (R < p) return 0; return (R-p)*t; } const bool operator<(const coupon &other) const { return t == other.t ? p < other.p : t < other.t; } }; vector<int> max_coupons(int A, vector<int> P, vector<int> T) { ll a = A, n = P.size(); ll required = 0, one = 0, two = 0, three = 0, four = 0; vector<coupon> coupons[5]; for (int i = 0; i < n; i++) { coupons[T[i]].push_back(coupon{T[i], P[i], i}), required += P[i]; if (T[i] == 1) one++; else if (T[i] == 2) two++; else if (T[i] == 3) three++; else four++; } for (int i = 1; i <= 4; i++) sort(all(coupons[i])); vector<ll> pref1; for (int i = 0; i < coupons[1].size(); i++) { if (i == 0) pref1.push_back(coupons[1][i].p); else pref1.push_back(pref1.back() + coupons[1][i].p); } int c2 = 0, c1 = 0; for (int i = 0; i < coupons[2].size(); i++) { int p = coupons[2][i].p; if (a - p < 0) break; if (a >= required) {c2 = two, c1 = one; break;} a = (a-p) * 2; int t1 = upper_bound(all(pref1), a) - pref1.begin(); if (i+t1+1 >= c2+c1) { c2 = i+1, c1 = t1; } } vector<int> ans; for (int i = 0; i < c2; i++) ans.push_back(coupons[2][i].idx); for (int i = 0; i < c1; i++) ans.push_back(coupons[1][i].idx); return ans; }