#include "festival.h" #include <bits/stdc++.h> using namespace std; vector<int> max_coupons(int A, vector<int> P, vector<int> T){ int sum = 0, n = P.size(); vector<pair<int, int> > c[3]; for(int i = 0; i < n; i++){ sum += P[i]; c[T[i]].push_back({P[i], i}); } for(auto &x : c) sort(x.begin(), x.end()); vector<int> pref_1 = {0}; for(auto x : c[1]) pref_1.push_back(pref_1.back() + x.first); int rem = A; pair<int, int> best = {upper_bound(pref_1.begin(), pref_1.end(), rem) - pref_1.begin() - 1, 0}; for(int i = 0; i < c[2].size(); i++){ rem = (rem - c[2][i].first) * 2; if(rem < 0) break; best = max(best, {i + 1 + upper_bound(pref_1.begin(), pref_1.end(), rem) - pref_1.begin() - 1, i + 1}); } vector<int> R; for(int i = 0; i < best.second; i++) R.push_back(c[2][i].second); while (R.size() < best.first) R.push_back(c[1][R.size() - best.second].second); return R; }