区間和問題 paiza問題集


はじめに

ここ数カ月、毎日こつこつ仕事終わりに勉強しています(主にpaizaスキルチェック、問題集)
そんな中で、今まであまり気にしていなかった累積和の考え方について自分用に記録する。

対象の問題

Paiza問題集「区間和 PHP編」

問題・条件・入力値

〇問題文
長さ N の数列 A と、K 個の区間 (l_1,r_1) ... (l_K,r_K) が与えられるので、各区間についての A の区間和 A_{l_i} + ... + A_{r_i} (1 ≦ i ≦ K) を求めてください。

・1 行目では、配列 A の要素数 N と与えられる整数の数 K が与えられます。
・続く N 行では、配列 A の要素が A_1 から順に与えられます。
・続く K 行では、和を求めるのに使う区間の値 l , r が与えられます。

〇条件
・1 ≦ N , K ≦ 100,000
・-100 ≦ A_i ≦ 100 (1 ≦ i ≦ N)
・1 ≦ l_i ≦ r_i ≦ N (1 ≦ i ≦ K)

//入力値
N K
A_1
...
A_N
l_1 r_1
...
l_K r_K

入力例と出力例

//入力例
4 2
16
88
10
-65
2 4
1 2
//出力例
33
104

実装

失敗例

この問題は、与えられた区間和を求める問題である。
初めて解いたとき、ただ指定の区間内でループ処理して区間和を求めればいいと思い以下のように実装しました。

list($N, $K) = explode(' ', trim(fgets(STDIN)));

//与えられる要素を配列に格納する
for ($i = 0; $i < $N; $i++) {
    $a[] = trim(fgets(STDIN));
}

for ($j = 0; $j < $K; $j++) {
    list($left, $right) = explode(' ', trim(fgets(STDIN)));
    $result = 0;
    //指定の区間内でループを行い、区間和を求める
    for ($k = $left-1; $k < $right; $k++) {
        $result += $a[$k];
    }
    echo $result."\n";
}

しかし処理時間が長く、タイムオーバーとなりました。
なぜ処理時間が長くなってしまったのか?

この実装では、指定区間が与えられたら毎回ループで算出しています。この実装では、計算量がO(N^2)となってしまいます。
今回の問題の条件では、1 ≦ N , K ≦ 100,000となっているため最大100000×100000もの計算量となってしまいタイムオーバーになってしまった。

解説と成功例

そこで、毎回指定の区間和を求めるのではなく、先に各区間の和を計算します。

以下に最初に示した入力例を図に表して考えてみる。

a0~a3が与えられる要素の値。sが各区間までの和である。
この図を基に計算してみると、各区間和は以下のように求められる。

s_{0} = 0
s_{1} = s_{0} + a_{0} = 0 + 16 = 16
s_{2} = s_{1} + a_{1} = 16 + 88 = 104
s_{3} = s_{2} + a_{2} = 104 + 10 = 114
s_{4} = s_{3} + a_{3} = 114 - 65 = 49

これらの式は以下のような関係となっている。

s_{i} = s_{i-1} + a_{i-1}

この関係を用いて、各区間の累積和sをあらかじめ計算することで計算量をO(N)に減らすことが出来る。

list($N, $K) = explode(' ', trim(fgets(STDIN)));

$a = []; //要素を格納する配列
$s = []; //各区間和を格納する配列
$s[0] = 0;

for ($i = 1; $i <= $N; $i++) {
    $a[] = trim(fgets(STDIN));
    $s[$i] = $s[$i-1] + $a[$i-1]; //上記で述べた関係式
}

for ($j = 0; $j < $K; $j++) {
    list($left, $right) = explode(' ', trim(fgets(STDIN)));
    $result = $s[$right] - $s[$left-1];
    echo $result."\n";
}

例えば、入力例の最初の指定区間である区間2~区間4までの和であればこのようになる。

そのため、

result = s_{4} - s_{1} = 49 - 16 = 33

となる。

図で考えてみるとイメージしやすい。

まとめ

今回は累積和について自分用にまとめた。
今後膨大なデータを扱う際に、このように処理時間を意識した実装を考えられるようになっていきたい。
これを機に様々なアルゴリズムに関しても少しずつ学んでいきたい。