#include "festival.h" #include <bits/stdc++.h> #define cost first #define ind second using namespace std; int c = 0; vector<int> pref1; 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) { for (int j = 0; j < options[0][i].first; ++j) pref1.push_back(c); ++c; } int t2 = options[1].size(), t3 = options[2].size(), t4 = options[3].size(); pair<int, 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; 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, dp[i][j][k].second}; 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, dp[i][j][k].second}; 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, dp[i][j][k].second}; dp[i][j][k+1].second.push_back(options[3][k].ind); } } } } } int a, b, o, d; vector<int> ret; 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; int n = i + j + k; if (dp[i][j][k].first >= pref1.size()) { ret = dp[i][j][k].second; for (int l = 0; l < c; ++l) { ret.push_back(options[0][l].ind); } } int _c = pref1[dp[i][j][k].first]; if (ret.size() < n + _c) { ret = dp[i][j][k].second; for (int l = 0; l < _c; ++l) { ret.push_back(options[0][l].ind); } } } } } return ret; }