/**
* __ __ __ __
* \ \ / / \ \ / (_) _____
* \ \ / /_ _\ \ / / _ ___|_ _|
* \ \/ /| | | |\ \/ / | |/ _ \ | |
* \ / | |_| | \ / | | __/ | |
* \/ \__,_| \/ |_|\____||_|
*
* Author: VuViet
* Created: 2025-08-09 19:44
**/
#include <bits/stdc++.h>
using namespace std;
const int N = 1e6 + 6;
struct Item
{
long long p;
int t, idx;
} items[N];
int n;
long long l_id, r_id;
bool cmp(Item a, Item b)
{
if (a.t == b.t && a.t == 1)
return a.p < b.p;
long long c1 = 1LL * a.p * a.t * (b.t - 1);
long long c2 = 1LL * b.p * b.t * (a.t - 1);
return c1 < c2;
}
vector<int> res, chosen[N / 5];
vector<long long> dp;
void DFS(long long pos, long long cnt)
{
if (pos < l_id || cnt == 0) return;
if (chosen[pos][cnt])
{
DFS(pos - 1, cnt - 1);
res.push_back(items[pos].idx);
} else {
DFS(pos - 1, cnt);
}
}
void Init(const vector<int> &P, const vector<int> &T)
{
n = (int)P.size();
for (int i = 0; i < n; i++) {
items[i] = { (long long)P[i], T[i], i };
}
sort(items, items + n, cmp);
r_id = n;
while (r_id > 0 && items[r_id - 1].t == 1)
r_id--;
l_id = r_id;
}
bool Greedy(int &A)
{
res.clear();
for (int i = 0; i < r_id; i++)
if ((A - items[i].p) * items[i].t >= A)
{
A = (int)((A - items[i].p) * items[i].t);
res.push_back(items[i].idx);
if (A > 1e16)
{
res.clear();
for (int j = 0; j < n; j++)
res.push_back(items[j].idx);
return true;
}
}
else
{
l_id = i;
break;
}
return false;
}
void DP(int A)
{
dp.clear();
dp.push_back(1LL * A);
for (long long i = l_id; i < r_id; i++)
{
int sz = (int)dp.size();
chosen[i].resize(sz, 0);
if ((dp[sz - 1] - items[i].p) * items[i].t >= 0)
{
dp.push_back((dp[sz - 1] - items[i].p) * items[i].t);
chosen[i].push_back(1);
}
for (int j = sz - 2; j >= 0; j--)
{
long long val = (dp[j] - items[i].p) * items[i].t;
if (val > dp[j + 1])
{
dp[j + 1] = val;
chosen[i][j + 1] = 1;
}
}
}
}
int FindBest()
{
int maxCount = 0, pos = 0;
for (int i = 0; i < (int)dp.size(); i++)
{
int cnt = 0;
long long lo = r_id, hi = n - 1;
while (lo < hi)
{
long long mid = (lo + hi + 1) / 2;
if (dp[i] >= items[mid].p)
{
cnt = (int)(mid - r_id + 1);
lo = mid;
}
else
{
hi = mid - 1;
}
}
if (i + cnt > maxCount)
{
maxCount = i + cnt;
pos = i;
}
}
return pos;
}
vector<int> max_coupons(int A, vector<int> P, vector<int> T)
{
Init(P, T);
if (Greedy(A)) return res;
for (int i = r_id + 1; i < n; i++) {
items[i].p += items[i - 1].p;
}
DP(A);
int pos = FindBest();
DFS(r_id - 1, pos);
long long remain = dp[pos];
for (int i = n - 1; i > r_id; i--) {
items[i].p -= items[i - 1].p;
}
for (int i = r_id; i < n; i++)
if (remain >= items[i].p)
{
remain -= items[i].p;
res.push_back(items[i].idx);
}
return res;
}
