#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
bool comp(array<long long,3>&a, array<long long,3>&b){
if(a[0]*a[1]*b[1]+b[0]*b[1]==b[0]*a[1]*b[1]+a[0]*a[1]){
return a[0]<b[0];
}
return a[0]*a[1]*b[1]+b[0]*b[1]<b[0]*a[1]*b[1]+a[0]*a[1];
}
vector<int> max_coupons(int a, vector<int> P, vector<int> T) {
long long A = a;
int n = P.size();
array<long long,3>arr[n];
int cn1 = 0;
for(int i = 0;i<n;i++){
arr[i]={P[i],T[i],i};
if(T[i]==1){
cn1++;
}
}
sort(arr,arr+n,comp);
if(cn1==0){
//only 2s here
for(int i = 0;i<n;i++){
A-=arr[i][0];
A*=arr[i][1];
if(A<0){
vector<int>ans;
for(int j = 0;j<i;j++){
ans.push_back(arr[j][2]);
}
return ans;
}
if(A>1e15){
vector<int>ans;
for(int i = 0;i<n;i++){
ans.push_back(arr[i][2]);
}
return ans;
}
}
vector<int>ans;
for(int i = 0;i<n;i++){
ans.push_back(arr[i][2]);
}
return ans;
}
//some 1s exist
long long prefsum[cn1];
for(int i = n-cn1;i<n;i++){
prefsum[i-(n-cn1)]=arr[i][0];
}
for(int i = 1;i<cn1;i++){
prefsum[i]+=prefsum[i-1];
}
int c = upper_bound(prefsum,prefsum+cn1,A)-prefsum;
pair<int,int> bestp = {c,c};
for(int i = 0;i<n-cn1;i++){
A-=arr[i][0];
A*=arr[i][1];
if(A<0){
break;
}
if(A>1e15){
pair<int,int>curr = {n,cn1};
bestp=curr;
break;
}
int c = (upper_bound(prefsum,prefsum+cn1,A)-prefsum);
pair<int,int>curr = {i+1+c,c};
bestp=max(bestp,curr);
}
int f = bestp.first-bestp.second;
vector<int>ans;
for(int i = 0;i<f;i++){
ans.push_back(arr[i][2]);
}
for(int i = 0;i<bestp.second;i++){
ans.push_back(arr[n-cn1+i][2]);
}
return ans;
}
