#include "festival.h"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
vector<int> max_coupons(int A, vector<int> P, vector<int> T) {
    int n = P.size();
    vector<vector<pair<int, int>>> can(4);
    for (int i = 0; i < n; i++) {
        T[i]--;
        can[T[i]].emplace_back(P[i], i);
    }
    for (int i = 0; i < 4; i++) {
        ranges::sort(can[i]);
    }
    vector<ll> pref(can[0].size() + 1, 0);
    for (int i = 0; i < can[0].size(); i++) {
        pref[i + 1] = pref[i] + can[0][i].first;
    }
    map<tuple<int, int, int>, pair<ll, int>> dp_cur, dp_next;
    map<tuple<int, int, int>, pair<ll, int>> saved;
    dp_cur[{0, 0, 0}] = {A, -1};
    bool has_inf = false;
    ll inf = 2'000'000'000'000'000;
    auto update_with_incr = [&](int i, int j, int l, int tr) -> void {
        ll res = dp_cur[{i, j, l}].first;
        if (tr == 1) {
            if (can[1].size() > i && (res - can[1][i].first) * 2 > res) {
                if ((res - can[1][i].first) * 2 > inf) has_inf = true;
                dp_next[{i + 1, j, l}] = max(dp_next[{i + 1, j, l}], {(res - can[1][i].first) * 2, 1});
            }
        } else if (tr == 2) {
            if (can[2].size() > j && (res - can[2][j].first) * 3 > res) {
                if ((res - can[2][j].first) * 3 > inf) has_inf = true;
                dp_next[{i, j + 1, l}] = max(dp_next[{i, j + 1, l}], {(res - can[2][j].first) * 3, 2});
            }
        } else {
            if (can[3].size() > l && (res - can[3][l].first) * 4 > res) {
                if ((res - can[3][l].first) * 4 > inf) has_inf = true;
                dp_next[{i, j, l + 1}] = max(dp_next[{i, j, l + 1}], {(res - can[3][l].first) * 4, 3});
            }
        }
    };
    auto update_with = [&](int i, int j, int l, int tr) -> void {
        ll res = dp_cur[{i, j, l}].first;
        if (tr == 1) {
            if (can[1].size() > i && res - can[1][i].first >= 0) {
                dp_next[{i + 1, j, l}] = max(dp_next[{i + 1, j, l}], {(res - can[1][i].first) * 2, 1});
            }
        } else if (tr == 2) {
            if (can[2].size() > j && res - can[2][j].first >= 0) {
                dp_next[{i, j + 1, l}] = max(dp_next[{i, j + 1, l}], {(res - can[2][j].first) * 3, 2});
            }
        } else {
            if (can[3].size() > l && res - can[3][l].first >= 0) {
                dp_next[{i, j, l + 1}] = max(dp_next[{i, j, l + 1}], {(res - can[3][l].first) * 4, 3});
            }
        }
    };
    auto max_ans = [&](int i, int j, int l) -> int {
        int l1 = 0, r1 = can[0].size();
        ll have = dp_cur[{i, j, l}].first;
        while (l1 != r1) {
            int m = (l1 + r1 + 1) / 2;
            if (pref[m] <= have) {
                l1 = m;
            } else {
                r1 = m - 1;
            }
        }
        return i + j + l + l1;
    };
    int now_sum = 0;
    while (true) {
        ll lst = 0;
        for (auto [k, v] : dp_cur) {
            for (int tr = 1; tr < 4; tr++) {
                update_with_incr(get<0>(k), get<1>(k), get<2>(k), tr);
            }
            saved[k] = v;
            lst = v.first;
        }
        if (dp_next.empty()) break;
        if (has_inf) break;
        swap(dp_cur, dp_next);
        dp_next.clear();
        now_sum++;
    }
    tuple<int, int, int> ans_coords = {-1, -1, -1};
    if (has_inf) {
        for (auto [k, v] : dp_next) {
            if (v.first > inf) {
                ans_coords = k;
                saved[k] = v;
                break;
            }
        }
    }
    if (!has_inf) {
        for (auto [k, v] : dp_cur) {
            ans_coords = k;
        }
        ll have = dp_cur[ans_coords].first;
        while (can[1].size() > get<0>(ans_coords) && (have - can[1][get<0>(ans_coords)].first) * 2 == have) {
            get<0>(ans_coords)++;
            saved[ans_coords] = {have, 1};
        }
        while (can[2].size() > get<1>(ans_coords) && (have - can[2][get<1>(ans_coords)].first) * 3 == have) {
            get<1>(ans_coords)++;
            saved[ans_coords] = {have, 2};
        }
        while (can[3].size() > get<2>(ans_coords) && (have - can[3][get<2>(ans_coords)].first) * 4 == have) {
            get<2>(ans_coords)++;
            saved[ans_coords] = {have, 3};
        }
        dp_cur.clear();
        dp_cur[ans_coords] = saved[ans_coords];
        // assert(dp_cur[ans_coords].first < 2'000'000'001 || get<0>(ans_coords) + get<1>(ans_coords) + get<2>(ans_coords) == can[1].size() + can[2].size() + can[3].size());
        int ans = max_ans(get<0>(ans_coords), get<1>(ans_coords), get<2>(ans_coords));
        int old_sum = now_sum;
        for (; now_sum < old_sum + 32; now_sum++) {
            for (auto [k, v] : dp_cur) {
                for (int tr = 1; tr < 4; tr++) {
                    update_with(get<0>(k), get<1>(k), get<2>(k), tr);
                }
                saved[k] = v;
            }
            swap(dp_cur, dp_next);
            dp_next.clear();
            for (auto [k, v] : dp_cur) {
                int ma = max_ans(get<0>(k), get<1>(k), get<2>(k));
                if (ma > ans) {
                    ans = ma;
                    ans_coords = k;
                }
            }
        }
        vector<int> vans;
        have = saved[ans_coords].first;
        while (ans_coords != make_tuple(0, 0, 0)) {
            int tr = saved[ans_coords].second;
            if (tr == 1) {
                vans.push_back(can[1][get<0>(ans_coords) - 1].second);
                get<0>(ans_coords)--;
            } else if (tr == 2) {
                vans.push_back(can[2][get<1>(ans_coords) - 1].second);
                get<1>(ans_coords)--;
            } else {
                vans.push_back(can[3][get<2>(ans_coords) - 1].second);
                get<2>(ans_coords)--;
            }
        }
        ranges::reverse(vans);
        int ptr = 0;
        while (ptr < can[0].size() && have >= can[0][ptr].first) {
            have -= can[0][ptr].first;
            vans.push_back(can[0][ptr].second);
            ptr++;
        }
        return vans;
    }
    vector<int> vans;
    tuple<int, int, int> s = ans_coords;
    while (ans_coords != make_tuple(0, 0, 0)) {
        int tr = saved[ans_coords].second;
        if (tr == 1) {
            vans.push_back(can[1][get<0>(ans_coords) - 1].second);
            get<0>(ans_coords)--;
        } else if (tr == 2) {
            vans.push_back(can[2][get<1>(ans_coords) - 1].second);
            get<1>(ans_coords)--;
        } else {
            vans.push_back(can[3][get<2>(ans_coords) - 1].second);
            get<2>(ans_coords)--;
        }
    }
    ranges::reverse(vans);
    for (int i = 0; i < can[0].size(); i++) {
        vans.push_back(can[0][i].second);
    }
    for (int i = get<0>(s); i < can[1].size(); i++) {
        vans.push_back(can[1][i].second);
    }
    for (int i = get<1>(s); i < can[2].size(); i++) {
        vans.push_back(can[2][i].second);
    }
    for (int i = get<2>(s); i < can[3].size(); i++) {
        vans.push_back(can[3][i].second);
    }
    return vans;
}
| # | 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... |