#include "festival.h" #include <bits/stdc++.h> #define cost first #define ind second #define ll long long using namespace std; ll c = 0; vector<ll> pref1{0}; ll bs(ll money) { ll l = 0, r = pref1.size(); while(r-l > 1) { // cout<<l<<' '<<r<<'\n'; ll m = (r+l)/2; if (pref1[m] <= money) { l = m; } else { r = m; } } return l; } std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) { vector<vector<pair<int, int>>> options(4, vector<pair<int, int>>{}); for (int i = 0; i < P.size(); ++i) { options[T[i]-1].push_back({P[i], i}); } for (int i = 0; i < 4; ++i) sort(options[i].begin(), options[i].end()); for (int i = 0; i < options[0].size(); ++i) { pref1.push_back(pref1.back() + options[0][i].first); ++c; } int t2 = options[1].size(), t3 = options[2].size(), t4 = options[3].size(); pair<ll, vector<int>> dp[t2+1][t3+1][t4+1]; for (int i = 0; i < t2; ++i) { for (int j = 0; j < t3; ++j) { for (int k = 0; k < t4; ++k) { dp[i][j][k] = {-1, vector<int>{}}; } } } dp[0][0][0] = {A, vector<int>{}}; for (int i = 0; i <= t2; ++i) { for (int j = 0; j <= t3; ++j) { for (int k = 0; k <= t4; ++k) { if (dp[i][j][k].first == -1) continue; vector<int> adding = dp[i][j][k].second; if (i != t2) { if ((dp[i][j][k].first - options[1][i].cost) * 2 > dp[i+1][j][k].first && options[1][i].cost <= dp[i][j][k].first) { dp[i+1][j][k] = {(dp[i][j][k].first - options[1][i].cost) * 2, adding}; dp[i+1][j][k].second.push_back(options[1][i].ind); } } if (j != t3) { if ((dp[i][j][k].first - options[2][j].cost) * 3 > dp[i][j+1][k].first && options[2][j].cost <= dp[i][j][k].first) { dp[i][j+1][k] = {(dp[i][j][k].first - options[2][j].cost) * 3, adding}; dp[i][j+1][k].second.push_back(options[2][j].ind); } } if (k != t4) { if ((dp[i][j][k].first - options[3][k].cost) * 4 > dp[i][j][k+1].first && options[3][k].cost <= dp[i][j][k].first) { dp[i][j][k+1] = {(dp[i][j][k].first - options[3][k].cost) * 4, adding}; dp[i][j][k+1].second.push_back(options[3][k].ind); } } } } } ll a, b, o, d; vector<int> ret; for (ll i = 0; i <= t2; ++i) { for (ll j = 0; j <= t3; ++j) { for (ll k = 0; k <= t4; ++k) { if (dp[i][j][k].first == -1) continue; ll n = i + j + k; if (dp[i][j][k].first >= pref1.back()) { ret = dp[i][j][k].second; for (ll l = 0; l < c; ++l) { ret.push_back(options[0][l].ind); } continue; } ll _c = bs(dp[i][j][k].first); if (ret.size() < n + _c) { ret = dp[i][j][k].second; for (ll l = 0; l < _c; ++l) { ret.push_back(options[0][l].ind); } } } } } return ret; }