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

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

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

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

問題概要

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