ABC233_D – Count Interval

問題
提出
1. 解けたが、もう少し簡単な方法があった

問題

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

提出

TLE: https://atcoder.jp/contests/abc233/submissions/44695634
AC: https://atcoder.jp/contests/abc233/submissions/44695692

解けたが、もう少し簡単な方法があった

以下、思いついた解法を書く。

累積和を用いる。
累積和$B_i$とすると、以下のようになる。
$$B_1 = A_1$$
$$B_2 = A_1+A_2$$
$$B_3 = A_1+A_2+A_3$$

つまり、$B_i$は$A_1$から$A_i$までの和になっている。
これをみると、$A_1$から連続する部分列の和がわかるので、Kになる部分を探す。
次に$A_2$から始まる部分については、$B_i-A_1$がKになる部分を探す。
次に$A_3$から始まる部分については、$B_i-A_1-A_2$がKになる部分を探す。
…
のようにすると答えが求められそう。
1から始まる部分を求めるとき、調べる必要がある$B_i$の個数は$N$個、2では$N-1$個、…となるので$O(N^2)$になる。
そこで、mapを使って$B_i$のうち、どの数字が何個あるかをカウントしておく。
これを使うことで$O(N)$で求められる。
mapを使うと、i番目以降の和を検索できなるので、全体の和からそれまでの和を引くことで対応した。

最初、mapでなくunordered_multisetを使って実装したがTLEになってしまった。
unordered_multisetのcountは、検索対象をaとして、$O(count(a))$かかる。
検索対象がt個だとt回かかる。
$B_i$が全て同じ和の時、$O(N^2)$になってしまう。

unordered_multiset::count - cpprefjp C++日本語リファレンス

キーを検索し、コンテナ内に見つかった要素の数を返す。

また、公式解説を見るともっと簡単に求めていた。
mapを使って和を検索するとき、自分はlを固定して計算していたが、公式解説ではrを固定して計算していた。
mapを使うと順番がわからなくなるので、lからrまでを検索する時、 lを固定して計算すると全体の和からi番目までのものを引くことになる。
しかし、rを固定するとr以内の累積和を探索するのみなので、全体からわざわざ引く必要はなくなる。