Official

D - Takahashi's Expectation Editorial by MMNMM


プレゼントを後ろから見ることでこの問題を解くことができます。

\(\operatorname{dp}\lbrack i\rbrack\lbrack j\rbrack\coloneqq i\) 番目のプレゼントを受け取った直後(\(i=0\) のとき、はじめのプレゼントを受け取る前)のテンションの値が \(j\) のときの、最終的なテンションの値 のような DP テーブルを考えます。 すると、\[\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}\] と計算できます。

しかし、\(0\le j\le10 ^ 9\) の範囲のすべての \(j\) に対して DP テーブルを計算していては実行時間制限に間に合わせるのは難しいです。 ここで、高橋くんのテンションの値 \(j\) がありうるプレゼントの価値 \(P _ i\) の範囲と比較して非常に大きくなっている場合、以降しばらくの高橋くんのテンションの変化が明らかに定まることを利用すると、DP テーブルの大きさを小さくすることができます。

適切な範囲で計算を行っておけば、DP テーブルで計算した範囲の外のクエリが与えられた場合、範囲内に入るまでテンションが下がり続けることから、範囲内に入るまでにプレゼントを受け取る個数を計算することで答えを求めることができます。

DP テーブルの大きさは、大きく \(2\) つの方針に分けられます。

  • \(\displaystyle0\le j\le\max _ iP _ i\) の範囲で計算する方針
    • メモリ使用量は小さくなりますが、前計算中に DP テーブルの外の値が必要になり、遷移が複雑になります。
  • \(\displaystyle0\le j\le\max _ i\lbrace A _ i+P _ i\rbrace\) の範囲で計算する方針
    • 前計算の段階で DP テーブルの外側の値が必要にならず、遷移はシンプルになりますが、メモリ使用量は大きくなります。

実装例は以下のようになります。 この実装例では、後者の方針で DP テーブルの計算を行っています。

#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;

    // テンションは max (P + A) 以下の値を管理する
    const auto M{ranges::max(P | views::transform([](const auto& p) { return p[0] + p[1]; }))};

    // dp[i][j] := i 番目までのプレゼントを受け取ったときのテンションが j のときの、最終的なテンションの値
    vector dp(N + 1, vector<unsigned>(M + 1));

    // すべて受け取り終わったときのテンションが j なら、最終的なテンションは 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, そうでなければ -B
            dp[i][j] = j <= p ? dp[i + 1][j + a] : dp[i + 1][j - min(j, b)];
    }

    // 大きめのテンションが来たときにいつ範囲内に入るかを判定するため、B の累積和を計算しておく
    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{});

    // はじめのテンションが x のときの最終的なテンションの値(オフセット込み)
    const auto access{[&dp, &sum_B, N, M](unsigned x) -> unsigned {
        // 最初から範囲内なら、dp の値を返せばよい
        if (x <= M)
            return dp[0][x];

        // 高いなら、しばらく下がり続ける
        const auto need_down{ranges::lower_bound(sum_B, x - M)};
        // N 回下がっても足りないなら、N 回下がるだけ
        if (need_down == end(sum_B))
            return x - sum_B.back();
        // 範囲に入ったら dp の値を返す
        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: