#include <bits/stdc++.h>
//#include "festival.h"
using namespace std;
long long MX=1e15;
long long fnd(vector <long long>& k, long long num){
    
    long long l=0;
    long long r=k.size()-1;
    if (num >= k[r]) return r;
    while (l<r) {
        long long md=(l+r+1)/2;
        if (k[md]<=num)
            l=md;     
        else
            r=md-1;  
    }
    return l;
}
vector<int> max_coupons(int A,vector<int> P,vector <int> T){
    
    vector <pair<long long,long long>> vp[5];
 
    for (int i=0; i<P.size(); i++){vp[T[i]].push_back({P[i],i});} 
    int sub_task_6=1;
    for (int i=0; i<P.size(); i++){if ((A-P[i])*T[i]>=A){sub_task_6=0;}}
    
    if (vp[2].size()==0 && vp[3].size()==0 && vp[4].size()==0){
    //subtask 1 / 5 points / T[i] = 1 for each i such that 0 ≤ i < N
    //for subtask 1 it is clear that we should sort and just do some o(n) solution.
    
    sort(vp[1].begin(),vp[1].end());vector <int> ans;int k=0;
    while (A>=0 && k<vp[1].size()){A=A-vp[1][k].first;ans.push_back(vp[1][k].second);k++;} 
    if (A<0)ans.pop_back();
    
    return (ans);}
    else if (vp[3].size()==0 && vp[4].size()==0){ 
    //subtask 2 / 7 points / N ≤ 3000; T[i] ≤ 2 for each i such that 0 ≤ i < N.
    //subtask 3 / 12 points /  T[i] ≤ 2 for each i such that 0 ≤ i < N. 
    
    //both of the subtasks have a common constraint T[i]<=2 hence we will use that to solve both the subtasks
    //Is there any point in 
    //  1: buying a T[i]==2 coupon after any T[i]==1. 
    //     because {X,Y|X is the P[i] where T[i]==1,Y is the P[i] where T[i]==2}
    //     ((A-X)*1-Y)*2<((A-Y)*2-X) where if we simplify it {0<X} hence no need.
    //     Inductively we can say that the structure should be something like (for the T values) 2,2,2,2,1,1,1,1,1
    //  2: for any structure if we seperate the 1s and 2s for the T values there can not be any {i,j|0<=i,j<=n-1,j=i+1}
    //     such that P[i]>=P[j] because ((A-P[i])*2-P[j])*2<((A-P[j])*2-P[i])*2 because if we simplify it we get
    //     P[j]<P[Pi] hence we must sort it.
    //  3: if at any point we have money>=2e14 we can just do anything we want since 2e14>=max(N)*max(P[i])
    //
    //With these two in mind for subtask 2 for every possible ending of the 2 chain
    //we use our subtask 1 solution to get the length of the max len of our ans if the 2 chain ended there 
    //hence our time complexity is o(vp2.size()*vp1.size())~o(n*n)=o(n^2)
    //
    //To get subtask 3 we just add Prefix sum followed by binary search to it.
    //hence our time complexity is o(vp2.size()*log(vp1.size()))~o(n*logn)
    vector <long long> psum1;
    vector <long long> psum2;
    long long sm1=0;
    long long sm2=A;
    
    vp[1].push_back({0,-1});
    vp[2].push_back({0,-1});
    sort(vp[1].begin(),vp[1].end());
    sort(vp[2].begin(),vp[2].end());
    
    for (long long i=0; i<vp[1].size(); i++){
        sm1+=vp[1][i].first;
        psum1.push_back(sm1);
    }
    
    for (long long i=0; i<vp[2].size(); i++){
        sm2-=vp[2][i].first;
        if (vp[2][i].first==0){
        psum2.push_back(sm2);            
        }
        else if (sm2>=0){
        sm2=sm2*2;
        sm2=min(sm2,MX);
        psum2.push_back(sm2);
        }
        else{
        psum2.push_back(-1);
        }
    }
    vector <int> ans;
    long long k=0,maxh=0,maxpos=-1;
    
    for (long long i=0; i<vp[2].size(); i++){
        if (psum2[i]>=0){
        long long kk=fnd(psum1,psum2[i]);
        if (kk+i>=maxh){maxpos=i;maxh=kk+i;}
    }}
    for (long long i=1; i<=maxpos; i++){
        ans.push_back(vp[2][i].second);
    }
    for (long long i=1; i<=maxh-maxpos; i++){
        ans.push_back(vp[1][i].second);        
    }
    return (ans);}
    //for online compiler use the /* */ around the <=70 subtask solution
    //because online compilers dont have that much space
    
    else if (P.size()<=70){
    //subtask 4 / 15 points / N ≤ 70    
    //it is obvious it is o(n^4) and we just have to make a 4d dp with with the info we got from 
    //the second and third subtask. One thing to note is that we dont really have to care abt the 
    //negative value until the end because the negative values wont turn positive ever.
    
    for (int i = 1; i <= 4; i++) {
        sort(vp[i].begin(), vp[i].end());
    }
    long long dp[71][71][71][71];
    for (int i=0; i<71; i++){
        for (int j=0; j<71; j++){
            for (int k=0; k<71; k++){
                for (int l=0; l<71; l++){dp[i][j][k][l]=-1;}}}}
                
    dp[0][0][0][0]=A;
    vector<int> vvv={0,0,0,0};
    
    int idx[5] = {0};
    
    for (idx[1]=0; idx[1]<=vp[1].size(); idx[1]++){
        for (idx[2]=0; idx[2]<=vp[2].size(); idx[2]++){
            for (idx[3]=0; idx[3]<=vp[3].size(); idx[3]++){
                for (idx[4]=0; idx[4]<=vp[4].size(); idx[4]++){
                    long long &tttt=dp[idx[1]][idx[2]][idx[3]][idx[4]];
                    for (int i=1; i<=4; i++) {
                        if (idx[i]==0)continue;
                        long long prv=dp[idx[1]-(i==1)][idx[2]-(i==2)][idx[3]-(i==3)][idx[4]-(i==4)];
                        tttt=max(tttt,(prv-vp[i][idx[i]-1].first)*i);
                    }
                    tttt=min(tttt,MX);
                    long long SM=0;
                    for (int i=0; i<vvv.size(); i++){SM+=vvv[i];}
                    
                    if(tttt>=0&&idx[1]+idx[2]+idx[3]+idx[4]>SM){vvv={idx[1],idx[2],idx[3],idx[4]};}
                }
            }
        }
    }    
    
    vector<int> ans;
    while (vvv[0]+vvv[1]+vvv[2]+vvv[3]){
        long long rem=dp[vvv[0]][vvv[1]][vvv[2]][vvv[3]];
        for (int i=1; i<=4; i++) {
            if (vvv[i-1]==0)continue;
            long long prv=dp[vvv[0]-(i==1)][vvv[1]-(i==2)][vvv[2]-(i==3)][vvv[3]-(i==4)];
            if (min((prv-vp[i][vvv[i-1]-1].first)*i,MX)==rem) {
                ans.push_back(vp[i][vvv[i-1]-1].second);
                vvv[i-1]--;
                break;
            }
        }
    }
    reverse(ans.begin(),ans.end());
    return ans;    
    
    
    
    
    
        
    }
    
    if (sub_task_6==1){
    //subtask 6 / 16 points / (A − P[i]) ⋅ T[i] < A for each i such that 0 ≤ i < N.
    //Even though it looks like this inequality wont give us much this is actually rapidly decreasing 
    //like those 1/x functions (IF WE REMOVE THE T[i]==1) becasue for a random time event where the current
    //money we have is A-x then (A-x-P[i])*T[i]<A-x*t[i] hence we can see it is decreasing rapidly
    //we can use the implementation we did on our subtask 2-3 the binary search.and the subtask 4
    
    for (int i=1; i<=4; i++){sort(vp[i].begin(), vp[i].end());}
    vector <long long> psum1;
    long long sm1=0;
    vp[1].push_back({0,-1});
    for (long long i=0; i<vp[1].size(); i++){
        sm1+=vp[1][i].first;
        psum1.push_back(sm1);
    }    
    
    long long dp[21][21][21];
    for (int j=0; j<21; j++){
        for (int k=0; k<21; k++){
            for (int l=0; l<21; l++){dp[j][k][l]=-1;}}}
                
    dp[0][0][0]=A;
    vector<long long> vvv={0,0,0};
    long long vvv1=fnd(psum1,0);
    int idx[5] = {0};
    
    for (idx[2]=0; idx[2]<=vp[2].size() && idx[2]<21; idx[2]++){
        for (idx[3]=0; idx[3]<=vp[3].size() && idx[3]<21; idx[3]++){
            for (idx[4]=0; idx[4]<=vp[4].size() && idx[4]<21; idx[4]++){
                long long &tttt=dp[idx[2]][idx[3]][idx[4]];
                for (int i=2; i<=4; i++) {
                    if (idx[i]==0)continue;
                    long long prv=dp[idx[2]-(i==2)][idx[3]-(i==3)][idx[4]-(i==4)];
                    tttt=max(tttt,(prv-vp[i][idx[i]-1].first)*i);
                }
                tttt=min(tttt,MX);
                long long SM=0;
                for (int i=0; i<vvv.size(); i++){SM+=vvv[i];}
                SM+=vvv1;
                
                if(tttt>=0&&fnd(psum1,tttt)+idx[2]+idx[3]+idx[4]>SM){vvv={idx[2],idx[3],idx[4]};vvv1={fnd(psum1,tttt)};}
            }
        }
    }
    
    
    vector<int> ans;
    
    for (int i=vvv1-1; i>=0; i--){
        ans.push_back(vp[1][i].second);
    }    
    //cout<<vvv[0]<<' '<<vvv[1]<<' '<<vvv[2]<<' '<<vvv[3]<<"\n";    
    
    while (vvv[0]+vvv[1]+vvv[2]){
        long long rem=dp[vvv[0]][vvv[1]][vvv[2]];
        
        for (int i=1; i<=3; i++) {
            if (vvv[i-1]==0)continue;
            
            long long prv=dp[vvv[0]-(i==1)][vvv[1]-(i==2)][vvv[2]-(i==3)];
            
            if (min((prv-vp[i+1][vvv[i-1]-1].first)*(i+1),MX)==rem) {
                ans.push_back(vp[i][vvv[i-1]-1].second);
                vvv[i-1]--;
                break;
            }
            //cout<<vvv1<<' '<<vvv[0]<<' '<<vvv[1]<<" "<<vvv[2]<<" "<<prv<<"\n";
        }
    }
    reverse(ans.begin(),ans.end());
    //cout<<"love";
    return ans;
    }
    else{
        
    //subtask 5 / 27 points / Nayra can buy all N coupons (in some order).
    //since there is a sequence where we can acquire every coupon we just have to find
    //a sequence that lets us have all of them hence we will no longer have to deal with some money issues
    //doing simple algebra gives the formula that if we pick (i,j) over (j,i) then 
    //P[i]*T[i]/(T[i]-1)<P[j]*T[j]/(T[j]-1) 
    long long a = A;
    vector<int> vv(P.size());
    for (int i=0; i<P.size(); i++)vv[i]=i;
    sort(vv.begin(), vv.end(), [&](int i, int j) {
        if (T[i]==T[j]) return P[i]<P[j];
        return 1ll*P[i]*T[i]*(T[j]-1)<1ll*P[j]*T[j]*(T[i]-1);
    });
    
    vector<int> ans;
    for(int i=0; i<P.size(); i++){
        if(a>=P[vv[i]]){
            a=(a-P[vv[i]])*T[vv[i]];
            a=min(a,MX);
            ans.push_back(vv[i]);
        }
    }
    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... |