#include "festival.h"
#include<utility>
#include<map>
#include<set>
#include<algorithm>
#include<queue>
#include<iostream>
const int nmax=2e5+42;
std::vector< std::pair<int,int> > ordered[5];
std::map< std::pair< std::pair<int,int>, int>, std::pair<int, long long> > scores;
long long prefix1[nmax];
std::queue< std::pair< std::pair<int,int>, int> > active;
void push_state(int diff,std::pair< std::pair<int,int>, int> new_state,long long new_score)
{
    active.push(new_state);
    if(scores.count(new_state)==0||scores[new_state].second<new_score)scores[new_state]={diff,new_score};
}
int eval_sign(long long a)
{
    if(a<0)return -1;
    if(a==0)return 0;
    return 1;
}
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
    long long total=0;
    int n=P.size();
    for(int i=0;i<n;i++)
    {
        total+=P[i];
        ordered[T[i]].push_back({P[i],i});
    }
    for(int i=1;i<=4;i++)
        sort(ordered[i].begin(),ordered[i].end());
    for(int i=0;i<ordered[1].size();i++)
    {
        if(i)prefix1[i]+=prefix1[i-1];
        prefix1[i]+=ordered[1][i].first;
    }
    scores[{{0,0},0}]={0,A};
    active.push({{0,0},0});
    std::pair< std::pair<int,int>, int> best_state={{0,0},0};
    int best_coupons=0;
    bool add_all=0;
    while(active.size())
    {
        std::pair< std::pair<int,int>, int> state=active.front();
        active.pop();
        if(scores[state].second==-1)continue;
        long long score=scores[state].second;
        scores[state].second=-1;
        if(score>=total)
        {
            add_all=1;
            best_state=state;
            break;
        }
        int pos2=state.first.first;
        int pos3=state.first.second;
        int pos4=state.second;
        int current_coupons=pos2+pos3+pos4;
        current_coupons+=std::upper_bound(prefix1,prefix1+ordered[1].size(),score)-prefix1;
        //std::cout<<pos2<<" "<<pos3<<" "<<pos4<<" -> "<<score<<" "<<current_coupons<<std::endl;
        if(current_coupons>best_coupons)
        {
            best_state=state;
            best_coupons=current_coupons;
        }
        int types[3]={-1,-1,-1};
        if(pos2<ordered[2].size()&&score>=ordered[2][pos2].first)types[0]=eval_sign((score-ordered[2][pos2].first)*2-score);
        if(pos3<ordered[3].size()&&score>=ordered[3][pos3].first)types[1]=eval_sign((score-ordered[3][pos3].first)*3-score);
        if(pos4<ordered[4].size()&&score>=ordered[4][pos4].first)types[2]=eval_sign((score-ordered[4][pos4].first)*4-score);
        int best_type=std::max(types[0],std::max(types[1],types[2]));
        if(pos2<ordered[2].size()&&score>=ordered[2][pos2].first&&best_type==types[0])
            push_state(2,{{pos2+1,pos3},pos4},(score-ordered[2][pos2].first)*2);
        if(pos3<ordered[3].size()&&score>=ordered[3][pos3].first&&best_type==types[1])
            push_state(3,{{pos2,pos3+1},pos4},(score-ordered[3][pos3].first)*3);
        if(pos4<ordered[4].size()&&score>=ordered[4][pos4].first&&best_type==types[2])
            push_state(4,{{pos2,pos3},pos4+1},(score-ordered[4][pos4].first)*4);
    }
    int best2=best_state.first.first;
    int best3=best_state.first.second;
    int best4=best_state.second;
    //std::cout<<best2<<" "<<best3<<" "<<best4<<" -> "<<best_coupons<<std::endl;
    std::vector<int> outp={};
    std::pair< std::pair<int,int>, int> initial={{0,0},0};
    while(best_state!=initial)
    {
        int cur=scores[best_state].first;
        int pos2=best_state.first.first;
        int pos3=best_state.first.second;
        int pos4=best_state.second;
        //std::cout<<best2<<" "<<best3<<" "<<best4<<" -> cur = "<<cur<<std::endl;
        if(cur==2)
        {
            pos2--;
            outp.push_back(ordered[2][pos2].second);
        }
        else if(cur==3)
        {
            pos3--;
            outp.push_back(ordered[3][pos3].second);
        }
        else
        {
            pos4--;
            outp.push_back(ordered[4][pos4].second);
        }
        best_state={{pos2,pos3},pos4};
    }
    std::reverse(outp.begin(),outp.end());
    if(add_all)
    {
        best_coupons=n;
        for(int i=best2;i<ordered[2].size();i++)
            outp.push_back(ordered[2][i].second);
        for(int i=best3;i<ordered[3].size();i++)
            outp.push_back(ordered[3][i].second);
        for(int i=best4;i<ordered[4].size();i++)
            outp.push_back(ordered[4][i].second);
    }
    for(int i=0;outp.size()<best_coupons;i++)
        outp.push_back(ordered[1][i].second);
    return outp;
}
| # | 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... |