Official

D - Takahashi's Expectation Editorial by en_translator


The problem can be solved by scanning the presents in reverse order.

Consider a DP (Dynamic Programming) table defined as follows: \(\operatorname{dp}\lbrack i\rbrack\lbrack j\rbrack\coloneqq \) the final mood value if the mood value right after receiving present \(i\) is (or if \(i=0\), if the initial mood value is) \(j\). Then, this can be computed as \[\operatorname{dp}\lbrack i\rbrack\lbrack j\rbrack=\begin{cases}\operatorname{dp}\lbrack i+1\rbrack\lbrack j+A _ i\rbrack&\quad(j\le P _ i)\\\operatorname{dp}\lbrack i+1\rbrack\lbrack\max\lbrace0,j-B _ i\rbrace\rbrack&\quad(\text{otherwise.})\end{cases}\]

However, if you evaluate the DP table for all \(0\le j\le10 ^ 9\), it is difficult to finish within the execution time limit. Here, notice that if Takahashi’s mood value \(j\) is far larger than possible value of a present \(P_i\), then the transition of his mood is mostly obvious; this helps us reducing the size of the DP table.

If we compute the DP table within an appropriate range, when you are given a query outside the computed range, the answer can be found by simply computing the number presents he receives until it reaches within the range, since his mood keep decreasing until that.

There are two strategies when deciding the size of the DP table.

  • Compute within \(\displaystyle0\le j\le\max _ iP _ i\).
    • The memory usage is small, but the precalculation requires values outside the DP table, making transitions complex.
  • Compute within \(\displaystyle0\le j\le\max _ i\lbrace A _ i+P _ i\rbrace\).
    • The precalculation does not require values outside the DP table, but the memory usage is large.

The sample code is below. In this sample code, we adopt the latter DP table range.

#include <iostream>
#include <vector>
#include <array>
#include <algorithm>
#include <numeric>
#include <ranges>

int main() {
    using namespace std;
    cin.tie(nullptr);
    ios::sync_with_stdio(false);

    unsigned N;
    cin >> N;
    vector<array<unsigned, 3>> P(N);
    for (auto&& [p, a, b] : P)
        cin >> p >> a >> b;

    // Manage mood up to {max (P + A)}
    const auto M{ranges::max(P | views::transform([](const auto& p) { return p[0] + p[1]; }))};

    // dp[i][j] := the final mood value if the mood value right after receiving present i is j.
    vector dp(N + 1, vector<unsigned>(M + 1));

    // If the mood after receiving all presents is j, then the final mood is j
    iota(begin(dp.back()), end(dp.back()), 0);

    for (unsigned i = N; i--; ) {
        const auto& [p, a, b] = P[i];
        for (unsigned j = 0; j <= M; ++j)
            // +A if the current mood value is not greater than the present's value, and -B otherwise
            dp[i][j] = j <= p ? dp[i + 1][j + a] : dp[i + 1][j - min(j, b)];
    }

    // Find the cumulative sums to determine when the mood falls into the range when a large mood value is found.
    vector<unsigned> sum_B(N);
    ranges::copy(P | views::elements<2>, begin(sum_B));
    inclusive_scan(begin(sum_B), end(sum_B), begin(sum_B), plus{});

    // The final mood value when the initial mood is `x` (including offset)
    const auto access{[&dp, &sum_B, N, M](unsigned x) -> unsigned {
        // If it is within the range initially, then directly return the dp value
        if (x <= M)
            return dp[0][x];

        // If larger, keep descending for a while
        const auto need_down{ranges::lower_bound(sum_B, x - M)};
        // If descending N times is not sufficient, just descend N times
        if (need_down == end(sum_B))
            return x - sum_B.back();
        // Once reaching the range, return the DP value
        return dp[need_down - begin(sum_B) + 1][x - min(x, *need_down)];
    }};

    unsigned Q;
    cin >> Q;
    for (unsigned i = 0; i < Q; ++i) {
        unsigned x;
        cin >> x;
        cout << access(x) << '\n';
    }
    cout << flush;
    return 0;
}

posted:
last update: