제출 #1316584

#제출 시각아이디문제언어결과실행 시간메모리
1316584lucas110550축제 (IOI25_festival)C++20
100 / 100
67 ms28592 KiB
#include <bits/stdc++.h>
using namespace std;

static inline long long sat_add(long long a, long long b, long long CAP) {
    __int128 v = (__int128)a + (__int128)b;
    if (v > CAP) return CAP;
    if (v < -(__int128)CAP) return -CAP; // not really needed here
    return (long long)v;
}

static inline long long sat_sub(long long a, long long b, long long CAP) {
    __int128 v = (__int128)a - (__int128)b;
    if (v > CAP) return CAP;
    if (v < -(__int128)CAP) return -CAP;
    return (long long)v;
}

static inline long long sat_mul(long long a, long long b, long long CAP) {
    __int128 v = (__int128)a * (__int128)b;
    if (v > CAP) return CAP;
    if (v < -(__int128)CAP) return -CAP;
    return (long long)v;
}

// Safe check for: token * (t-1) >= p * t with saturation not needed for correctness
static inline bool feasible_prefix(long long token, int p, int t) {
    __int128 L = (__int128)token * (t - 1);
    __int128 R = (__int128)p * t;
    return L >= R;
}

std::vector<int> max_coupons(int A, std::vector<int> P, std::vector<int> T) {
    int n = (int)P.size();

    // coeff[t] = 12*t/(t-1) for t in {2,3,4}
    const int coeff[5] = {0, 0, 24, 18, 16};

    // ---------- split coupons ----------
    vector<pair<long long,int>> cheap; // (price, idx) for T=1

    struct Booster {
        long long key; // p * coeff[t] (fits in 64-bit for p<=1e9)
        int p;
        int t;
        int idx;
    };
    vector<Booster> boosters;
    cheap.reserve(n);
    boosters.reserve(n);

    long long SUM_CHEAP = 0;
    int MAXP = 0;

    for (int i = 0; i < n; i++) {
        int ti = T[i];
        int pi = P[i];
        MAXP = max(MAXP, pi);
        if (ti == 1) {
            cheap.push_back({(long long)pi, i});
            SUM_CHEAP += (long long)pi;
        } else {
            long long key = 1LL * pi * coeff[ti];
            boosters.push_back({key, pi, ti, i});
        }
    }

    // Choose a saturation cap large enough that exact values above it don't matter.
    // K=61 from the algorithm; extra slack is fine.
    const int K = 61;
    long long CAP = SUM_CHEAP;
    CAP = sat_add(CAP, sat_mul(4LL * MAXP, (long long)(K + 5), (long long)9e18), (long long)9e18);
    CAP = sat_add(CAP, A, (long long)9e18);

    // ---------- sort ----------
    sort(cheap.begin(), cheap.end(),
         [](const auto& a, const auto& b){ return a.first < b.first; });

    sort(boosters.begin(), boosters.end(),
         [](const Booster& a, const Booster& b){
            if (a.key != b.key) return a.key < b.key;
            return a.p < b.p;
         });

    // cheap prefix sums (64-bit safe)
    vector<long long> cheap_psum(cheap.size() + 1, 0);
    vector<int> cheap_idx(cheap.size());
    for (size_t i = 0; i < cheap.size(); i++) {
        cheap_psum[i + 1] = cheap_psum[i] + cheap[i].first;
        cheap_idx[i] = cheap[i].second;
    }

    auto cheap_cnt_from_token = [&](long long tok) -> int {
        // If tok is saturated beyond SUM_CHEAP, we can buy all cheap.
        if (tok >= SUM_CHEAP) return (int)cheap.size();
        auto it = upper_bound(cheap_psum.begin(), cheap_psum.end(), tok);
        return int(it - cheap_psum.begin()) - 1;
    };

    // ---------- take all non-decreasing boosters ----------
    long long token = min<long long>(A, CAP);
    vector<int> prefix_ids;
    prefix_ids.reserve(boosters.size());

    int i = 0, m = (int)boosters.size();
    while (i < m) {
        const auto &b = boosters[i];
        if (feasible_prefix(token, b.p, b.t)) {
            prefix_ids.push_back(b.idx);
            // token = (token - p) * t  with saturation
            token = sat_mul(sat_sub(token, b.p, CAP), b.t, CAP);
            i++;
        } else break;
    }

    // ---------- hard suffix DP ----------
    int h_len = m - i;
    vector<long long> dp(K + 1, -1);
    dp[0] = token;

    vector<vector<uint8_t>> parent(h_len, vector<uint8_t>(K + 1, 0));

    for (int j = 0; j < h_len; j++) {
        const auto &b = boosters[i + j];
        int max_c = min(j + 1, K);
        for (int c = max_c; c >= 1; c--) {
            long long prev = dp[c - 1];
            if (prev < 0) continue;
            if (prev >= b.p) {
                long long cand = sat_mul(sat_sub(prev, b.p, CAP), b.t, CAP);
                if (cand > dp[c]) {
                    dp[c] = cand;
                    parent[j][c] = 1;
                }
            }
        }
    }

    // ---------- evaluate all possibilities ----------
    int best_total = -1;
    int best_c = 0;
    long long best_token = dp[0];

    for (int c = 0; c <= min(K, h_len); c++) {
        long long cur = dp[c];
        if (cur < 0) continue;
        int cheap_cnt = cheap_cnt_from_token(cur);
        int total = i + c + cheap_cnt;
        if (total > best_total) {
            best_total = total;
            best_c = c;
            best_token = cur;
        }
    }

    // ---------- reconstruct chosen hard subset ----------
    vector<int> chosen_hard_pos;
    chosen_hard_pos.reserve(best_c);

    int need = best_c;
    for (int j = h_len - 1; j >= 0 && need > 0; j--) {
        if (parent[j][need]) {
            chosen_hard_pos.push_back(j);
            need--;
        }
    }
    reverse(chosen_hard_pos.begin(), chosen_hard_pos.end());

    // ---------- build final answer ----------
    vector<int> answer;
    answer.reserve(prefix_ids.size() + chosen_hard_pos.size() + cheap.size());

    for (int id : prefix_ids) answer.push_back(id);
    for (int pos : chosen_hard_pos) answer.push_back(boosters[i + pos].idx);

    int cheap_cnt_final = cheap_cnt_from_token(best_token);
    for (int k = 0; k < cheap_cnt_final; k++) answer.push_back(cheap_idx[k]);

    return answer;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...