#include "festival.h"
#include<bits/stdc++.h>
#define pb push_back
using namespace std;
const long long maxn = 2e5 + 10;
long long n, a;
struct coupon
{
long long cost, type, index;
coupon() {};
coupon(long long _cost, long long _type, long long _index)
{
cost = _cost;
type = _type;
index = _index;
}
};
bool cmp(coupon c1, coupon c2)
{
if(c1.type == c2.type)return (c1.cost < c2.cost);
return (c2.cost * c1.type * c2.type + c1.cost * c1.type > c1.type * c1.cost * c2.type + c2.cost * c2.type);
}
bool cmp2(coupon c1, coupon c2)
{
return (c1.cost < c2.cost);
}
vector < coupon > g;
vector < coupon > ones;
vector < long long > pref;
int getones(long long tokens)
{
int l = 0, r = pref.size()-1, mid, ans = 0;
while(l<= r)
{
mid = (l + r)/2;
if(pref[mid] <= tokens)
{
ans = mid;
l = mid + 1;
}
else r = mid - 1;
}
return ans;
}
long long gothrough(long long tokens, int index)
{
return (tokens - g[index].cost) * g[index].type;
}
long long dp[maxn][61], taken[maxn][61], from[maxn][61];
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T)
{
a = A;
n = P.size();
for (long long i = 0; i < n; ++ i)
{
if(T[i] != 1)g.pb(coupon(P[i], T[i], i));
else ones.pb(coupon(P[i], T[i], i));
}
sort(ones.begin(), ones.end(), cmp2);
pref.pb(0);
long long curr = 0;
for (auto &[c, t, i]:ones)
{
curr += c;
pref.pb(curr);
}
sort(g.begin(), g.end(), cmp);
vector < int > res;
int prefix = -1;
for (auto &[c, t, i]: g)
{
if(a >= 1LL * 1e17)
{
res.pb(i);
prefix ++;
continue;
}
long long aft = (a - c) * t;
if(aft >= a)
{
a = aft;
prefix ++;
res.pb(i);
}
else break;
}
if(prefix == g.size()-1)
{
int cnt1 = getones(a);
for (int i = 0; i < cnt1; ++ i)
res.pb(ones[i].index);
}
else
{
memset(dp, -1, sizeof(dp));
dp[prefix+1][0] = a;
dp[prefix+1][1] = gothrough(a, prefix+1);
taken[prefix+1][1] = 1;
int sz = g.size()-1;
for (int i = prefix+1; i < sz; ++ i)
{
for (int j = 0; j <= 60; ++ j)
{
if(dp[i][j] == -1)continue;
if(dp[i+1][j] < dp[i][j])
{
dp[i+1][j] = dp[i][j];
taken[i+1][j] = 0;
}
/// we take the next one
long long newtokens = gothrough(dp[i][j], i+1);
if(j < 60 && dp[i+1][j+1] < newtokens && newtokens >= 0)
{
dp[i+1][j+1] = newtokens;
taken[i+1][j+1] = 1;
}
}
}
int mx = 0, best = 0, cnt1 = 0;
for (int i = 0; i <= 60; ++ i)
{
if(dp[sz][i] == -1)continue;
long long curr = i + getones(dp[sz][i]);
if(curr > mx)
{
mx = curr;
best = i;
cnt1 = getones(dp[sz][i]);
}
}
vector < int > p;
int i = sz, j = best;
while(i > prefix)
{
if(taken[i][j])
{
p.pb(i);
j --;
}
i --;
}
reverse(p.begin(), p.end());
for (auto x: p)
res.pb(g[x].index);
for (int i = 0; i < cnt1; ++ i)
res.pb(ones[i].index);
}
return res;
}
