#include "festival.h" using namespace std; #include <bits/stdc++.h> static int how_many_type_1_can_we_buy(long long A2, const vector<long long>& type1_prefix) { // type1_prefix[0] = 0, increasing int k = int(upper_bound(type1_prefix.begin(), type1_prefix.end(), A2) - type1_prefix.begin()) - 1; return max(0, k); } std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) { __int128 tokens = A; struct Coupon { long long cost; int type; int index; bool operator<(const Coupon& other) const { if (cost != other.cost) return cost < other.cost; return type > other.type; // your tie-break } }; vector<Coupon> type2, type1; for (int i = 0; i < (int)P.size(); i++) { if (T[i] == 2) type2.push_back({P[i], T[i], i}); else type1.push_back({P[i], T[i], i}); } sort(type2.begin(), type2.end()); sort(type1.begin(), type1.end()); vector<long long> type1_prefix(1, 0); type1_prefix.reserve(type1.size() + 1); for (auto &c : type1) { type1_prefix.push_back(type1_prefix.back() + c.cost); } // baseline: buy only type1 long long A_ll = (long long)min<__int128>(tokens, LLONG_MAX); int best = how_many_type_1_can_we_buy(A_ll, type1_prefix); int best_t2 = 0, best_t1 = best; __int128 cur = tokens; for (int i = 0; i < (int)type2.size(); i++) { if (cur < type2[i].cost) break; cur -= type2[i].cost; cur *= 2; if(cur > LONG_LONG_MAX){ best = T.size(); best_t1 = type1.size(); best_t2 = type2.size(); break; } long long cur_ll = (long long)min<__int128>(cur, LLONG_MAX); int can1 = how_many_type_1_can_we_buy(cur_ll, type1_prefix); if (i + 1 + can1 > best) { best = i + 1 + can1; best_t2 = i + 1; best_t1 = can1; } } vector<int> res; res.reserve(best_t2 + best_t1); for (int i = 0; i < best_t2; i++) res.push_back(type2[i].index); for (int i = 0; i < best_t1; i++) res.push_back(type1[i].index); return res; }