Educational Codeforces Round 69 D: Yet Another Subarray Problem

問題概要

長さ $N$ の整数列 $a$ と、パラメータ $m, k$ が与えられる。

数列内の連続する閉区間 $a \left[ l, r \right]$ について、そのコストは $\sum_{i=l}^{r} a_i - k \lceil \frac{r-l+1}{m} \rceil$ で表される。なお、空の区間に対するコストは $0$ である。

あり得る全ての区間を考慮した時のコストの最大値を求めよ。

解法

$m$ で割った切り上げ、というパートが存在しなければよくある問題で、最適な左端のみを覚えながら右端を $1$ ずつ伸ばしていけば良い (累積和を取って変形すると、左端に関しては累積和の $\min$ が必要になる)。

しかし今回は切り上げがあるので少し工夫が必要。

インデックスを $\bmod m$ で考え、合同となるインデックスはまとめて管理しよう。 $i$ 番目までインデックスを伸ばすときは $(i-1) \bmod m$ における最適値が変わる可能性があり、直前の値を持ってくるか、それまで最適だったものをそのまま使うかの $2$ 通りを考慮する必要がある。

ソースコード

int main() {
    ll N, M, K; scanf("%lld%lld%lld", &N, &M, &K);
    vector<ll> A(N);
    for(int i=0; i<N; i++) scanf("%lld", &A[i]);
    vector<ll> S(N+1);
    for(int i=0; i<N; i++) S[i+1] = S[i] + A[i];

    ll ans = 0;
    vector<ll> opt(M, INF);
    for(int i=1; i<=N; i++) {
        opt[(i-1)%M] = min(opt[(i-1)%M] + K, S[i-1] + K);
        for(int k=0; k<M; k++) chmax(ans, S[i] - opt[k]);
    }
    printf("%lld\n", ans);
    return 0;
}

本番では $m$ が大きい場合しか解けなかった (右端を $i$ まで伸ばしたときに $k$ について何倍するか変化するのは高々 $\frac{i}{m}$ 個なので、変化するところだけ愚直に更新していっても $m$ が $\sqrt{N}$ 以上だと全体でも $O(N \sqrt{N})$ なので使えるはず)