精進記録

言語はc++です

D - Candy Distribution

こちらの問題です。

atcoder.jp

AtCoder Tagsのmap問題を埋めていてこの問題が難しかったので記事にしておきます。

問題概要

与えられた大きさNの数列Aのある範囲の総和を考えます。
この総和が自然数Mの倍数である範囲の取り方の個数を求めてください。

というような問題です。

制約

1≦N≦105
2≦M≦109
1≦A_i≦109

解法

愚直に全ての部分和を考えられれば一番楽ですが、その場合、O(n2)となるため間に合いません。累積和を考えます。

詰まります

今回のようなABC-D問題に挑戦するような人であれば累積和までの発想はそう難しくはないと思います。が、そこまでで自分は詰まりました。

参考URL

こういう時はだいたい「アルゴリズム けんちょん」で検索すると助かることが多いです。今回も「累積和 けんちょん」で検索しました。(けんちょんさんいつもありがとうございます。)

qiita.com

std::mapを使うことによって累積和問題の実装をより簡単にしています。

助かりませんでした

mapを使うところまではいいのですが、その先の実装までたどり着くことが出来ませんでした解説を読みます。

数列 A に対し、A の先頭 i 項の和を Bi とします。なお、便宜上 B0 = 0 としておきます。

いま、Al +...+Ar = Br −Bl−1 であり、Al +...+Ar が M の倍数であることは Br, Bl−1 を M で割ったあまりが等しいことと同値です。 よって、数えるべきは、Bi, Bj が mod M で等しいような 0 ≤ i < j ≤ Nの組の個数です。

解説では累積和配列をBとしています。
部分和が累積和配列の項の差で求められるのだから、部分和がMの倍数であるかどうかについては、累積和配列の各項のmod Mを求めて一致しているかを調べればよい、というわけです。

実装

        //入力
        int n,m; cin>>n>>m;
        vector<long> a(n);
        for(int i=0;i<n;i++)cin>>a[i];

        //累積和
        vector<long> b(n+1,0);
        for(int i=0;i<n;i++)b[i+1]=b[i]+a[i];

        //累積和の各項のmodをとる
        vector<long> mod(n+1);
        for(int i=0;i<n+1;i++)mod[i]=b[i]%m;

これでmodが取れました。
入力例2を入れると下記のようになります。

13 17 //入力
29 7 5 7 9 51 7 13 8 55 42 9 81 //入力
0 29 36 41 48 57 108 115 128 136 191 233 242 323 //累積和の出力
0 12 2 7 14 6 6 13 9 0 4 12 4 0 //累積和のmod 17の出力

累積和のmodの配列第2項に注目します。
第2項は12です。また、mod 17が12となる項はもう1つありますので、こことの差を取ってやることで17の倍数となる部分和が作れそうです。ans++できます。

mod 17が0となる場合にはどうでしょうか。3つあります。
第1、10、14項がそれにあたります。この時、1と10、10と14、14と1の累積和の差を取ってやることで部分和が作れることになります。ans+=3です。

このように共通のmodの個数を数えてその個数からansに足していけばよいということが分かります。
いくつ足すかについてはこれはnC2つまりn*(n-1)/2を足してやればよいです。
これをstd::mapを使うことで実装します。

        map<long,long> M; //小文字mだと入力と被る
        for(int i=0;i<n+1;i++)M[mod[i]]++;
        
        long ans=0;
        for(auto it : M)ans+=it.second*(it.second-1)/2;
        cout<<ans<<endl;

ここまでのコードをくってつければACします。
mapはかなり実用的なコンテナだとは思いますが、現状あまりスマートに使えた覚えがないのでもう少し勉強をしたいと思います。

ここまでお読みいただいた方ありがとうござました。