#include <bits/stdc++.h>
using namespace std;
using i128 = __int128_t;
struct Coupon {
int idx;
long long P;
int T;
};
static bool cmpMul(const Coupon& a, const Coupon& b) {
// Compare a.P * a.T / (a.T-1) < b.P * b.T / (b.T-1)
// cross multiply to avoid division:
// a.P * a.T * (b.T-1) < b.P * b.T * (a.T-1)
i128 left = (i128)a.P * (i128)a.T * (i128)(b.T - 1);
i128 right = (i128)b.P * (i128)b.T * (i128)(a.T - 1);
return left < right;
}
static bool cmpOne(const Coupon& a, const Coupon& b) {
return a.P < b.P;
}
std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
int N = (int)P.size();
vector<Coupon> mul, one;
mul.reserve(N);
one.reserve(N);
for (int i = 0; i < N; ++i) {
Coupon c{ i, (long long)P[i], T[i] };
if (T[i] == 1) one.push_back(c);
else mul.push_back(c);
}
// sort multiplier coupons by ascending C = P*T/(T-1)
sort(mul.begin(), mul.end(), cmpMul);
// sort T=1 coupons by price ascending
sort(one.begin(), one.end(), cmpOne);
int M = (int)mul.size();
const i128 INF = (i128)1 << 120; // huge sentinel
// DP[i][k] = max tokens after considering first i mul-coupons and taking exactly k of them
vector<vector<i128>> dp(M + 1, vector<i128>(M + 1, -1));
vector<vector<char>> took(M + 1, vector<char>(M + 1, 0));
dp[0][0] = (i128)A;
for (int i = 1; i <= M; ++i) {
const Coupon& c = mul[i - 1];
for (int k = 0; k <= i; ++k) {
// skip case
if (dp[i - 1][k] != -1) {
dp[i][k] = dp[i - 1][k];
took[i][k] = 0;
}
// take case
if (k > 0 && dp[i - 1][k - 1] != -1) {
i128 prev = dp[i - 1][k - 1];
if (prev >= (i128)c.P) {
i128 after = (prev - (i128)c.P) * (i128)c.T;
if (after > INF) after = INF;
if (after > dp[i][k]) {
dp[i][k] = after;
took[i][k] = 1;
}
}
}
}
}
// prefix sums for T=1 coupons
int N1 = (int)one.size();
vector<i128> pref(N1 + 1, 0);
for (int i = 1; i <= N1; ++i) {
pref[i] = pref[i - 1] + (i128)one[i - 1].P;
}
int bestK = 0, bestOneCnt = 0, bestTotal = 0;
for (int k = 0; k <= M; ++k) {
if (dp[M][k] == -1) continue;
i128 tokens = dp[M][k];
// find max number of T=1 coupons we can afford
int lo = 0, hi = N1;
while (lo < hi) {
int mid = (lo + hi + 1) >> 1;
if (pref[mid] <= tokens) lo = mid;
else hi = mid - 1;
}
int cntOne = lo;
int total = k + cntOne;
if (total > bestTotal) {
bestTotal = total;
bestK = k;
bestOneCnt = cntOne;
}
}
// reconstruct chosen multiplier coupons
vector<int> orderMul;
int i = M, k = bestK;
while (i > 0) {
if (took[i][k]) {
orderMul.push_back(mul[i - 1].idx);
--k;
}
--i;
}
reverse(orderMul.begin(), orderMul.end());
// take first bestOneCnt cheap T=1 coupons
vector<int> orderOne;
for (int idx = 0; idx < bestOneCnt; ++idx) {
orderOne.push_back(one[idx].idx);
}
// concatenate results
vector<int> result;
result.reserve(bestTotal);
result.insert(result.end(), orderMul.begin(), orderMul.end());
result.insert(result.end(), orderOne.begin(), orderOne.end());
return result;
}
# | 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... |