【競プロ】累積和を覚える

競プロやっていると、そのうちアルゴリズムの勉強が必要になってきます。今回はその中でも有名な累積和を扱います。

そこまで難しいものでもないので、すぐに実践で使えるようになるはずです。

1 累積和とは
- 1.1 配列の準備
- 1.2 任意の区間の和を求める
2 実践
3 終わりに

累積和とは

累積和は「数列のある区間の和を高速に計算するアルゴリズム」です。

「数列 \(a_n\) の第 \(n\) 項までの和を \(S_n\) とする。」なんて数学でやりますが、まさにこれと同じです。

例えば \(\sum_{k=15}^{50}a_n\) を求める時、ループを回して計算してもいいですが、求める値は \(S_{50}-S_{14}\) とも表せます。前者は \(35\) 回程度の計算が必要ですが、後者は \(1\) 回の計算で済みます。（これについてはまた後述します。）

配列の準備

では、実装の話をしましょう。

数列 \(a_0,a_1,\cdots ,a_n\) を配列 \(A\) に入れておきます。\(A[0]=a_0,A[1]=a_1,\cdots\) です。

配列 \(S\) を考えます。\(S[i]\) は「\(A\)の先頭から \(i\) 個の要素の和」とします。「数列の第 \(i\) 項までの和」とも言えますね。（数列は第 \(1\) 項, 第 \(2\) 項 \(\cdots\) のように数えてください。）

ただし、\(S[0]=0\) とします。(前者の定義で考えると自明。)

つまり、\(S\) は次のように計算しておけばいいでしょう。（\(S\) の要素数は、\(A\)の要素数\(+1\) だけ必要です。）

vector<int> S(n+1);
for(int i = 0;i < n;i++) S[i+1] = S[i]+A[i];

vectorを使う場合は初期値が \(0\) になっているのでこれでいいですが、静的配列を使う場合は \(S[0] = 0\) を忘れないようにしましょう。

任意の区間の和を求める

\(A[l]\) から \(A[r]\) の和を「区間 \([l,r]\) の和」のように書くとします。これは \((l-1,r+1)\) のように開区間でも書けるとします。

上のように \(S\) を準備すると、\([i,j)\ (=[i,j-1])\) の和は \(S[j]-S[i]\) と書けることが分かります。

例を挙げれば、

\([5,23]\) の和は \(S[24]-S[5]\)
\([0,i]\)の和は \(S[i+1]-S[0]\)
数列の第 \(2\) 項から第 \(7\) 項までの和は \(S[7]-S[1]\)

特に数列の第 \(i\) 項の番号と、配列 \(A\) のインデックスはずれるので注意しましょう。（このような場合はあまりないですが。）

実践

例題です。

以下が提出コードです。重りの番号は \(0,1,\cdots\) として考えると混乱しにくいです。

#include <bits/stdc++.h>
#define rep(i,n) for(int i = 0;i < (n);i++)
using namespace std;

int main(){
    int n;
    cin >> n;
    vector<int> w(n);
    rep(i,n) cin >> w[i];
    vector<int> s(n+1);
    rep(i,n) s[i+1] = s[i]+w[i];
    int ans = 100000;
    rep(t,n-1){
        int s1 = s[t+1];
        int s2 = s[n] - s[t+1];
        ans = min(ans,abs(s1-s2));
    }
    cout << ans << endl;
    return 0;
}

問題文の通りシミュレーションすれば解けます。

区間の和を取るところでループを回して計算すると、\(\mathcal{O}(N^2)\) になりますが、累積和を使えば \(\mathcal{O}(N)\) となります。

\(S\) の計算で \(\mathcal{O}(N)\)、区間の和の計算で \(\mathcal{O}(1)\) ですね。これが累積和を使うと嬉しいところです。