累積和を使う代表的な問題ですね。
問題概要
個の整数からなる数列 と、整数 が与えられる。
数列から連続する 個の値を選んで総和をとる。この総和として考えられる最大値を求めよ。
制約
前提知識
まず、愚直な解法では 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); }
この解法で一見正解できそうに思えるのですが、計算量は となります。これでは時間オーバーとなります。計算量という概念については、次の記事を参考にしてください。
計算量を改善するために、累積和という概念を活用します。
考えたこと
ここまでの前提知識を学べば、この問題は解けます。
「数列の連続する部分の総和」は累積和を求めることで管理できます。今回も、 の累積和 を次のように定義しましょう。
- ...
この を用いると の区間 の総和は
というふうに簡潔に表せます。よって今回の問題は
- を動かしていったときの
- の値の最大値
を求めれば OK です。計算量は となり、実行制限時間内に十分間に合います。
コード
#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; }