#include "festival.h"
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define ll long long
#define ld long double
#define ull unsigned long long
#define ff first
#define ss second
#define pii pair<int,int>
#define pll pair<long long, long long>
#define vi vector<int>
#define vl vector<long long>
#define pb push_back
#define rep(i, b) for(int i = 0; i < (b); ++i)
#define rep2(i,a,b) for(int i = a; i <= (b); ++i)
#define rep3(i,a,b,c) for(int i = a; i <= (b); i+=c)
#define count_bits(x) __builtin_popcountll((x))
#define all(x) (x).begin(),(x).end()
#define siz(x) (int)(x).size()
#define forall(it,x) for(auto& it:(x))
using namespace __gnu_pbds;
using namespace std;
typedef tree<int, null_type, less<int>, rb_tree_tag,tree_order_statistics_node_update> ordered_set;
//mt19937 mt;void random_start(){mt.seed(chrono::time_point_cast<chrono::milliseconds>(chrono::high_resolution_clock::now()).time_since_epoch().count());}
//ll los(ll a, ll b) {return a + (mt() % (b-a+1));}
const int INF = 1e9+50;
const ll INF_L = 1e18+40;
const ll MOD = 1e9+7;
ll P[200001];
ll T[200001];
bool comp(int a, int b)
{
    ll val_a = -P[a]*T[a]*T[b]-P[b]*T[b];
    ll val_b = -P[b]*T[b]*T[a]-P[a]*T[a];
    if(val_a != val_b) return val_a > val_b;
    if(T[a] != T[b]) return T[a] > T[b];
    return P[a] < P[b];
}
pll dp[200002][52];
vi max_coupons(int A2, vi P2, vi T2) 
{
    ll coins = A2;
    int n = siz(P2);
    rep(i,n)
    {
        P[i] = P2[i];
        T[i] = T2[i];
    }
    vi v;
    rep(i,n) v.pb(i);
    sort(all(v),comp);
    vi ans;
    int start_ind = n;
    rep(i,n)
    {
        if(coins >= 1e15)
        {
            rep2(j,i,n-1)
            {
                ans.pb(v[j]);
            }
            return ans;
        }
        if((coins-P[v[i]]) * T[v[i]] >= coins)
        {
            ans.pb(v[i]);
            coins = (coins-P[v[i]])*T[v[i]];
        }
        else
        {
            start_ind = i;
            break;
        }
    }
    if(start_ind == n) return ans;
    vi v2;
    vi ones;
    rep2(i,start_ind,n-1)
    {
        if(T[v[i]] > 1) v2.pb(v[i]);
        else ones.pb(v[i]);
    }
    rep(i,siz(v2)+1)
    {
        rep(d,51)
        {
            dp[i][d] = {-1,-1};
        }
    }
    dp[0][0] = {coins,-1};
    rep(i,siz(v2))
    {
        rep(d,51)
        {
            dp[i+1][d] = max(dp[i+1][d],dp[i][d]);
            if((dp[i][d].ff-P[v2[i]])*T[v2[i]] >= 0)
            {
                dp[i+1][d+1] = max(dp[i+1][d+1],{(dp[i][d].ff-P[v2[i]])*T[v2[i]],i});
            }
        }
    }
    pii best = {0,0};
    rep(d,51)
    {
        if(dp[siz(v2)][d].ff == -1) continue;
        ll c = dp[siz(v2)][d].ff;
        int ans2 = d;
        forall(it,ones)
        {
            if((c-P[it])*T[it] >= 0)
            {
                ans2++;
                c = (c-P[it])*T[it];
            }
        }
        best = max(best,{ans2,d});
    }
    int opt = best.ss;
    int cur_elm = siz(v2);
    vi ans2;
    while(dp[cur_elm][opt].ss != -1)
    {
        ans2.pb(v2[dp[cur_elm][opt].ss]);
        cur_elm = dp[cur_elm][opt].ss;
        opt--;
    }
    reverse(all(ans2));
    forall(it,ans2) ans.pb(it);
    coins = A2;
    forall(it,ans) coins = (coins-P[it])*T[it];
    forall(it,ones) if((coins-P[it])*T[it] >= 0)
    {
        coins = (coins-P[it])*T[it];
        ans.pb(it);
    }
    return ans;
}
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... |