#include "festival.h" using namespace std; #include <bits/stdc++.h> static int how_many_type_1_can_we_buy(long long A2, const vector<long long>& type1_prefix) { // type1_prefix[0] = 0, increasing int k = int(upper_bound(type1_prefix.begin(), type1_prefix.end(), A2) - type1_prefix.begin()) - 1; return max(0, k); } std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) { __int128 tokens = A; struct Coupon { long long cost; int type; int index; bool operator<(const Coupon& other) const { if (cost != other.cost) return cost < other.cost; return type > other.type; // your tie-break } }; vector<Coupon> type4, type3, type2, type1; for (int i = 0; i < (int)P.size(); i++) { if(T[i] == 4){ type4.push_back({P[i], T[i], i}); } else if(T[i] == 3) { type3.push_back({P[i], T[i], i}); } else if (T[i] == 2){ type2.push_back({P[i], T[i], i}); } else { type1.push_back({P[i], T[i], i}); } } sort(type4.begin(), type4.end()); sort(type3.begin(), type3.end()); sort(type2.begin(), type2.end()); sort(type1.begin(), type1.end()); vector<long long> type1_prefix(1, 0); type1_prefix.reserve(type1.size() + 1); for (auto &c : type1) { type1_prefix.push_back(type1_prefix.back() + c.cost); } // dp[i][j][k] = max tokens after using i type 2, j type 3, k type 4 struct DPstate { __int128 val; array<int, 3> prev; int item_index; }; vector<vector<vector<DPstate>>> dp(type2.size()+1, vector<vector<DPstate>>(type3.size()+1, vector<DPstate>(type4.size()+1, {0, {0, 0, 0}, -1}))); int best_numitems = 0; array<int, 3> bestDPstate = {0, 0, 0}; int best_type1 = 0; dp[0][0][0].val = A; for(int i = 0; i <= type2.size(); i ++){ for(int j = 0; j <= type3.size(); j ++){ for(int k = 0; k <= type4.size(); k ++){ if(i > 0){ DPstate new_state = dp[i-1][j][k]; new_state.val -= type2[i-1].cost; new_state.val *= 2; new_state.prev = {i-1, j, k}; new_state.item_index = type2[i-1].index; if(new_state.val > dp[i][j][k].val && new_state.val >= 0){ dp[i][j][k] = new_state; } } if(j > 0){ DPstate new_state = dp[i][j-1][k]; new_state.val -= type3[j-1].cost; new_state.val *= 3; new_state.prev = {i, j-1, k}; new_state.item_index = type3[j-1].index; if(new_state.val > dp[i][j][k].val && new_state.val >= 0){ dp[i][j][k] = new_state; } } if(k > 0){ DPstate new_state = dp[i][j][k-1]; new_state.val -= type4[k-1].cost; new_state.val *= 4; new_state.prev = {i, j, k-1}; new_state.item_index = type4[k-1].index; if(new_state.val > dp[i][j][k].val && new_state.val >= 0){ dp[i][j][k] = new_state; } } if(dp[i][j][k].item_index == -1){ //dp[i][j][k].val = 0; //continue; } int items_bought = i + j + k; int also_type1 = how_many_type_1_can_we_buy(dp[i][j][k].val, type1_prefix); //cout << also_type1; if(items_bought + also_type1 > best_numitems && dp[i][j][k].item_index != -1){ best_numitems = items_bought + also_type1; bestDPstate = {i, j, k}; best_type1 = also_type1; } } } } //cout << best_numitems; vector<int> res; array<int, 3> cur = bestDPstate; if(dp[cur[0]][cur[1]][cur[2]].item_index == -1) return {}; while(1){ res.push_back(dp[cur[0]][cur[1]][cur[2]].item_index); cur = dp[cur[0]][cur[1]][cur[2]].prev; if(cur[0] == 0 && cur[1] == 0 && cur[2] == 0){ break; } } reverse(res.begin(), res.end()); for(int i = 0; i < best_type1; i ++){ if(i < type1.size()){ res.push_back(type1[i].index); } } return res; } /* 12 10 1 2 3 2 2 2 6 2 6 2 1 2 2 2 5 1 5 1 5 2 20 1 25 1 3 1 2 4 5 3 7 1 State = [how many type 2 we have used][how many type 3 we have used] dp[i][j] = max( (dp[i-1][j] - type2[i].price) * 2, (dp[i][j-1] - type3[i].price) * 3 ) */