部分配列の和 = k の個数

何本の区間を、何回足すことになるのか

連続する部分配列のうち、和が k になるものが何本あるかを数えます。まず、区間そのものが何本あるかを数えてみます。

区間を数えるだけならこれで済みますが、1 本ごとに和を出す必要があります。左端を決めて右へ伸ばしながら足していけば、1 本あたりの足し算は 1 回で済みます。それでも 1 万個の配列で 5000 万回です。

同じ足し算を、何度もやり直している

[3, 4, 7, 2] で、左端が先頭の区間と、1 つ右にずらした区間を並べてみます。

• 左端が 3 — 3、3+4、3+4+7、3+4+7+2
• 左端が 4 — 4、4+7、4+7+2

4+7 は両方に出てきます。左端をずらすたびに、その右側の足し算をもう一度やり直しているわけです。

やり直しを消す方法があります。先頭からの和を 1 本だけ持っておきます。

4 と 7 の区間の和は 14 - 3 で 11 です。どの区間の和も、右端までの和から左端の手前までの和を引くだけで出ます。足し算はもう終わっていて、残るのは引き算 1 回だけになりました。

1 周で数えるために、何を覚えておけばよいか

右端を 1 つずつ右へずらしながら考えます。右端までの和を S とすると、和が k になる区間の左端は、S - k という和で終わっている位置です。

つまり毎回の問いは「S - k という和が、これまでに何回出てきたか」になります。前回は「出たかどうか」だけで足りましたが、今回は回数が要ります。区間が何本あるかを数える問題だからです。

覚えておくのは、これまでに出た先頭からの和と、それぞれが何回出たか。右端を 1 つ進めるたびに 1 回問い合わせて、返ってきた回数をそのまま答えに足していきます。配列を歩くのは 1 周だけです。

まだ何も足していない状態も、1 回に数える

[3, 4] で k = 7 のときを追ってください。右端が 4 の位置で S は 7、S - k は 0 です。和が 0 で終わっている位置は、まだ 1 つも足していない先頭の状態しかありません。ここを数に入れておかないと、配列の先頭から始まる区間が丸ごと数え落とされます。

もう 1 つ、和が k を超えたら左端を縮める、という書き方は今回使えません。負の数が混ざっていると、右へ伸ばした先で和が減ることがあるからです。先頭からの和で考えるかぎり、負の数も 0 もそのまま扱えます。

1. 累積和 + hashmap を使い O(n) で実装する
2. 配列に負の数や 0 が含まれても正しく動作すること
3. 戻り値は部分配列の個数 (int)