Codeforces round 627 Div.3 E

March 16, 2020

https://codeforces.com/contest/1324/problem/E

解いたのでメモ。

以下、法を $h$ として考える。

1~n番目の睡眠それぞれで、 $a_{i}$ を選ぶか $a_{i} - 1$ を選ぶか決められるので、選び方の組み合わせは全部で $2^n$ 個ある。 $1 \le n \le 2000$ なので、全探索では制限時間に間に合わない。そこで、i番目の睡眠をしたとき、 $j(\le i)$ 番目の睡眠時間の選択でどう選んだかという情報は、「現在時刻」に集約されることに着目して動的計画法で解く。

$dp[i + 1][j] := i\text{番目の睡眠をしたときに時刻がjであるとき「いい睡眠」をした回数の最大値}$

とすると、jが到達可能な状態か、つまりi番目の睡眠をしたときに時刻jになっていることがありうるかどうかを考慮しないといけない。これだと少しややこしいので、

$m_{i} := \text{1~}i\text{番目までの睡眠で}a_j - 1 (1\le j \le i)\text{を選択した回数の総和}$

と定義すると、i番目の睡眠を終えた時点で、現在時刻 $t$ が、

$t = \sum_{j = 1}^{i}a_j - m_i$

となることから、

$dp[i + 1][m] := i\text{番目までの睡眠で}m\text{回}a_j - 1\,\,(1\le j \le i)\text{を選んだとき「いい睡眠」をした回数の最大値}$

と定義して問題を解く。状態 $(i, m)$ のとき、その状態の時刻を求めるためには、 $\sum_{j = 1}^{i}a_{j}$ を求める必要があるが、事前に累積和を求めておくことで、 $O(1)$ で計算できる。状態 $(i, m)$ からは、 $a_i$ を選んで $(i + 1, m)$ へ遷移する場合と、 $a_i - 1$ を選んで $(i + 1, m + 1)$ へ遷移する場合がある。

実装

全体の計算量は $O(N^2)$

// テンプレ省略

#define MAX_N 2000
#define MAX_H 2000
i64 N, H, L, R;
i64 A[MAX_N];
i64 dp[MAX_N + 1][MAX_N + 1];

int main() {
    cin >> N >> H >> L >> R;
    REP(i, N) cin >> A[i];

    vector<i64> C(N + 1, 0);
    REP(i, N) {
        C[i + 1] = C[i] + A[i];
        C[i + 1] %= H;
    }

    REP(i, N + 1) {
        REP(j, N + 1) {
            dp[i][j] = 0;
        }
    }

    i64 nt;
    REP(i, N) {
        for (i64 j = 0; j <= i; j++) {
            REP(k, 2) {
                nt = C[i + 1] - ((j + k) % H);
                if (nt < 0) nt += H;
                chmax(dp[i + 1][j + k], dp[i][j] + (L <= nt && nt <= R));
            }
        }
    }

    i64 max_n = 0;
    REP(i, N + 1) {
        chmax(max_n, dp[N][i]);
    }
    cout << max_n << endl;

    return 0;
}