累積和を使う代表的な問題ですね。

問題概要

$N$ 個の整数からなる数列 $A_{0}, \dots, A_{N-1}$ と、整数 $K$ が与えられる。

数列から連続する $K$ 個の値を選んで総和をとる。この総和として考えられる最大値を求めよ。

制約

$1 \le K \le N \le 10^{5}$

前提知識

まず、愚直な解法では TLE になることを確認しましょう。たとえば、次のように考えたとします。

int res = 0;
for (int i = 0; i + K < N; ++i) {
    // A[i] から連続する K 個の値を求める
    int sum = 0;
    for (int j = i; j < i + K; ++j) {
        sum += A[i];
    }

    // 最大値を更新する
    res = max(res, sum);
}

この解法で一見正解できそうに思えるのですが、計算量は $O(NK)$ となります。これでは時間オーバーとなります。計算量という概念については、次の記事を参考にしてください。

qiita.com

計算量を改善するために、累積和という概念を活用します。

qiita.com

考えたこと

ここまでの前提知識を学べば、この問題は解けます。

「数列の連続する部分の総和」は累積和を求めることで管理できます。今回も、 $A$ の累積和 $S$ を次のように定義しましょう。

$S_{0} = 0$
$S_{1} = A_{0}$
$S_{2} = A_{0} + A_{1}$
$S_{3} = A_{0} + A_{1} + A_{2}$
...
$S_{N} = A_{0} + A_{1} + A_{2} + \dots + A_{N-1}$

この $S$ を用いると $A$ の区間 $\lbrack l, r)$ の総和は

$A_{l} + A_{l+1} + \dots + A_{r-1} = S_{r} - S_{l}$

というふうに簡潔に表せます。よって今回の問題は

$i$ を動かしていったときの
$S_{i+K} - S_{i}$ の値の最大値

を求めれば OK です。計算量は $O(N)$ となり、実行制限時間内に十分間に合います。

コード

#include <iostream>
#include <vector>
using namespace std;
const long long INF = 1LL<<60; // 仮想的な無限大の値

int main() {
    // 入力
    int N, K; cin >> N >> K;
    vector<long long> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
        
    // 累積和を前計算
    vector<int> s(N+1, 0);
    for (int i = 0; i < N; ++i) s[i+1] = s[i] + a[i];
    
    // 答えを求める
    long long res = -INF; // 最初は無限小の値に初期化しておく
    for (int i = 0; i <= N-K; ++i) {
        long long val = s[K+i] - s[i];
        if (res < val) res = val;
    }
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

JOI 本選 2007 A - 最大の和 (AOJ 0516, 難易度 4)

問題概要

制約

前提知識

考えたこと

コード