#include<bits/stdc++.h> #include "festival.h" using namespace std; typedef long long ll; struct coupon { int p, t, i; ll eval(ll x) { return (x-p)*t; } bool operator<(coupon j) { if (t==1) return j.p > p; return j.eval(eval(0)) > eval(j.eval(0)); } }; std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) { ll n=P.size(), tot_p=0; deque<coupon> ord, one; vector<ll> pref; for (int i = 0; i < n; i++) { if (T[i]==1) one.push_back({P[i], 1, i}); else ord.push_back({P[i], T[i], i}); tot_p+=P[i]; } sort(ord.begin(), ord.end()); if (one.size()) { sort(one.begin(), one.end()); pref.push_back(one[0].p); for (int i = 1; i < one.size(); i++) { pref.push_back(pref[i-1]+one[i].p); } } // greedy select the coupons that increase tokens taking care of overflow ll X=A; vector<int> R; while (true) { if (!ord.empty() && X<ord[0].eval(X)) { R.push_back(ord[0].i); X=ord[0].eval(X); tot_p-=ord[0].p; ord.pop_front(); n--; } else break; if (X>=tot_p) { for (auto i : ord) R.push_back(i.i); for (auto i : one) R.push_back(i.i); return R; } } // dp[num of coupons bought][pos in coupon sequence]=max possible number of tokens n=ord.size(); vector<vector<pair<ll, int>>> dp; dp.push_back(vector<pair<ll, int>>(n+1, {X, -1})); for (int i = 1; i <= n; i++) { dp.push_back(vector<pair<ll, int>>(n+1, {-4e18, -1})); for (int j = i; j <= n; j++) { ll buy=ord[j-1].eval(dp[i-1][j-1].first); if (buy>dp[i][j-1].first) dp[i][j]={buy, j-1}; else dp[i][j]=dp[i][j-1]; } if (dp[i][n].first<0) break; } // find optimal number of non-ones to buy int m=dp.size(); int sz=0, opt=0; for (int i = 0; i < m; i++) { int X=dp[i][n].first; if (dp[i][n].first<0) break; int ones=upper_bound(pref.begin(), pref.end(), X)-pref.begin(); if (ones+i>sz) { opt=i; sz=ones+i; } } sz-=opt; // recover sequence stack<int> S; int j=n; while (opt>0 && j>0) { int t=dp[opt][j].second; S.push(ord[t].i); opt--; j=t; } while (!S.empty()) { R.push_back(S.top()); S.pop(); } for (int i = 0; i < sz; i++) { R.push_back(one[i].i); } return R; }