#include "festival.h"
#include <algorithm>
#include <climits>
long long Mult(long long a, int t)
{
if (LLONG_MAX / t <= a)
return LLONG_MAX - 1;
return a * t;
}
long long Step(std::vector<std::pair<long long, int>> &P, int &j, long long a, int c)
{
long long a2 = LLONG_MAX;
if (j < P.size() && P[j].first <= a)
{
a2 = Mult(a - P[j].first, c);
if (a2 < a)
a2 = LLONG_MAX;
}
return a2;
}
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T)
{
int n = P.size();
long long a = A;
std::vector<std::pair<long long, int>> x[4];
for (int i = 0; i < 4; i++)
x[i] = std::vector<std::pair<long long, int>>();
for (int i = 0; i < n; i++)
{
x[T[i] - 1].push_back(std::make_pair(1LL * P[i], i));
}
for (int i = 0; i < 4; i++)
{
std::sort(x[i].begin(), x[i].end());
}
std::vector<int> ans = {};
int ind[] = {0, 0, 0, 0};
while (true)
{
long long a2 = Step(x[1], ind[1], a, 2);
long long a3 = Step(x[2], ind[2], a, 3);
long long a4 = Step(x[3], ind[3], a, 4);
long long a5 = std::min(std::min(a2, a3), a4);
if (a5 == LLONG_MAX)
break;
if (a5 == a2)
{
ans.push_back(x[1][ind[1]].second);
a = a2;
ind[1]++;
continue;
}
if (a5 == a3)
{
ans.push_back(x[2][ind[2]].second);
a = a3;
ind[2]++;
continue;
}
{
ans.push_back(x[3][ind[3]].second);
a = a4;
ind[3]++;
}
}
if (x[2].size() > 0 || x[3].size() > 0)
{
for (int j = 3; j >= 0; j--)
{
for (int i = ind[j]; i < x[j].size(); i++)
{
if (x[j][i].first <= a)
{
ans.push_back(x[j][i].second);
a = Mult(a - x[j][i].first, j + 1);
}
}
}
}
else
{
for (int i = 1; i < x[0].size(); i++)
x[0][i].first += x[0][i - 1].first;
int max = 0, mj = -1, mi = -1;
{
int l = -1, r = x[0].size();
while (l + 1 < r)
{
int m = (l + r) >> 1;
if (x[0][m].first <= a)
l = m;
else
r = m;
}
max = l + 1;
mj = ind[1] - 1;
mi = l;
}
for (int j = ind[1]; j < x[1].size(); j++)
{
if (x[1][j].first <= a)
{
a -= x[1][j].first;
a = Mult(a, 2);
int l = -1, r = x[0].size();
while (l + 1 < r)
{
int m = (l + r) >> 1;
if (x[0][m].first <= a)
l = m;
else
r = m;
}
if (j - ind[1] + l + 2 > max)
{
max = j - ind[1] + l + 2;
mj = j;
mi = l;
}
}
else
{
break;
}
}
for (int i = ind[1]; i <= mj; i++)
{
ans.push_back(x[1][i].second);
}
for (int i = 0; i <= mi; i++)
{
ans.push_back(x[0][i].second);
}
}
return ans;
}
