累積和の練習

問題概要

$N$ 個の整数 $A_1, A_2, \dots, A_N$ が与えられる。

各 $k = 1, 2, \dots, N$ に対して、数列の $k$ 個の連続する区間の値の和の最大値を求めよ。

制約

$1 \le N \le 3000$

考えたこと

愚直にやると計算量は

$k$ の値が $O(N)$ 通り考える必要があって
各 $k$ に対して $O(N)$ 個の区間があって
各区間に対して $O(N)$ 個の値を合計する必要があって

というわけで合計で $O(N^{3})$ の計算量がかかってしまう。ここで「区間の合計値」を参照する機会が何回もあることに気付く。それを高速に実現する手段として累積和と呼ばれる典型テクニックがある。配列 A に対して累積和とは

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] + ... + A[N-1]

で定義される配列 S のことである。ここで配列 A の範囲が [0, N) なら、配列 S の範囲が [0, N+1) になっていることに注意する。この配列 S は簡単に作ることができて、

S[0] = 0;
for (int i = 0; i < N; ++i) S[i+1] = S[i] + A[i];

みたいにすれば OK。 $O(N)$ で構成できる。さて、この累積和を用いると以下のことが言える！

数列 A の区間 [left, right) の和が S[right] - S[left] で与えられる

つまり、累積和を予め構築しておけば、任意の区間和が $O(1)$ で高速に求められることを意味する。

ここで、区間 [left, right) は左側が閉区間、右側が開区間になっていることに注意する。C++ の STL や Python のリストなど、多くの標準ライブラリが区間をこのような形式で扱っている。

元の問題に戻り、

$k$ の値が $O(N)$ 通り考える必要があって
各 $k$ に対して $O(N)$ 個の区間があって
各区間に対して $O(N)$ 個の値を合計する必要があって

のうちの最後のパートについては累積和を求めておくことで $O(1)$ で実現できる。具体的には S[i+k] - S[i] を各 i について計算して最大値をとればよい。

#include <iostream>
#include <vector>
using namespace std;

int main() {
    int N; cin >> N;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    // 累積和を作る
    vector<long long> S(N + 1, 0);
    for (int i = 0; i < N; ++i) S[i + 1] = S[i] + A[i];

    // 各 k に対して
    for (int k = 1; k <= N; ++k) {
        long long res = 0;
        for (int i = 0; i + k <= N; ++i) res = max(res, S[i + k] - S[i]);
        cout << res << endl;
    }
}