#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)); } }; bool DEBUG=false; std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) { // exchange argument on coupon order 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()); sort(one.begin(), one.end()); if (one.size()) { pref.push_back(one[0].p); for (int i = 1; i < one.size(); i++) { pref.push_back(pref[i-1]+one[i].p); } } if (DEBUG) { cerr << "order\n"; for (auto i : ord) { cerr << i.i << ' '; } cerr << '\n'; for (auto i : one) { cerr << i.i << ' '; } cerr << '\n'; for (auto i : pref) { cerr << i << ' '; } cerr << '\n'; } // 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; } } if (DEBUG){ cerr << "greedy\n"; for (auto i : R) { cerr << i << ' '; } cerr << '\n'; for (auto i : ord) { cerr << i.i << ' '; } cerr << '\n'; } // 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; } if (DEBUG) { cerr << "dp\n"; for (auto i : dp) { int idx=0; for (auto j : i) { cerr << '(' << idx++ << ')' << j.first << '-' << j.second << ' '; } cerr << "\n====================\n"; } cerr << '\n'; } // find optimal number of non-ones to buy int m=dp.size(); if (DEBUG) cerr << m << ' ' << n << '\n'; 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; if (pref.size()) ones=upper_bound(pref.begin(), pref.end(), X)-pref.begin(); else ones=0; if (DEBUG) cerr << "in loop: " << i << ' ' << ones << ' ' << X << ' ' << pref.size() << '\n'; if (ones+i>sz) { opt=i; sz=ones+i; } } sz-=opt; if (DEBUG) cerr << "out " << sz << ' ' << opt << '\n'; // recover sequence stack<int> S; int j=n; if (DEBUG) cerr << ord.size() << '\n'; while (opt>0 && j>0) { int t=dp[opt][j].second; if (DEBUG) cerr << opt << '-' << j << '-' << t << ' '; S.push(ord[t].i); opt--; j=t; } while (!S.empty()) { R.push_back(S.top()); S.pop(); } if (DEBUG) cerr << '\n'; if (DEBUG) cerr << sz << ' ' << R.size() << '\n'; for (int i = 0; i < sz; i++) { R.push_back(one[i].i); } if (DEBUG) cerr << "recovery\n"; return R; }