累積和(AGC 023 A問題)

今まで書いてきた競プロカテゴリの記事は自分で読み返すためだけのメモ書き程度の扱いでしたが、今回は他人に読まれることを意識して書いてみます。

累積和とは、初項からある項までの和のことで、区間の中の和を効率よく求めるアルゴリズムでよく使われるらしいです。先日開催されたAtcoder Grand Contest 023のA問題はこれを知らないと詰む問題だったようです。具体的にこの問題を使って累積和を勉強します。

問題文の概要は、

「整数列Aの部分列で、その総和が0になるものの個数を求めよ。」

というものです。

もし全パターンの部分和を愚直に計算しようとすると、forループを三重に回してO(N^3)になりそうです。しかし累積和を使うとこれをO(N^2)で求められます。

例えば、A={1,3,-4,2,2,-2}とすると、各項までの累積和は下図のように表せます。

ここで例えば第2項「3」から第5項「2」までの和を求めたければ、第5項までの累積和から第1項までの累積和を引けば良いですね。


このことから、部分和が0になるというのは、累積和の差が0、すなわち累積和が等しい二ヵ所を選べばよいと分かります。

あとはこのような選び方の組み合わせをどのようにして求めるか、ですが、累積和の数列を作ってソートしたら分かりやすいでしょうか。
上の累積和をソートすると
{0,0,1,2,2,4,4}
となり、等しい二項の組み合わせは
(0,0),(2,2),(4,4)の三通りとなります。

他にもmapを使ってもうまく数えられそうです。(僕はこちらを選びました)

最後に自分のコードを載せておきます。

1. //nCrを計算する関数
   long nCr(int n,int r) {
2. long ans = 1;
3. for (int i = n; i > n - r; i--) {
4. ans *= i;
5. }
6. for (int i = 1; i <= r; i++) {
7. ans /= i;
8. }
9. return ans;
10. }
11.  
12. int n;
    //累積和を入れる配列
13. long long c[200002];
14. int main() {
15. cin >> n;
    //第0項までの累積和を0とする
16. c[0] = 0;
    //累積和の値に対するValueを1増やす
17. map<long long, int>mp;
18. mp[c[0]]++;
19. for (int i = 0; i < n; i++) {
20. long long a;
21. cin >> a;
22. c[i + 1] = c[i] + a;
    //累積和の値に対するValueを1増やす
23. mp[c[i + 1]]++;
24. }
25. long long ans = 0;
    //mpの要素についてループする
26. for (auto a : mp) {
27. int p = a.second;
28. if (p >= 2) {
29. ans += nCr(p, 2);
30. }
31. }
32. cout << ans << endl;
33. }