ABC105-D Candy Distribution

AtCoder ABC 400点 累積和 数え上げ 易しめ

atcoder.jp


問題概要

N個の箱が左右一列に並んでいて、左からi番目の箱にはAi個のお菓子が入っている。連続した箱[l,r]からM人の子供に平等に配れるようにMの倍数個取り出したい時、l,rの選び方の総数を求めよ。

制約

  • 1 <= N <= 1e5
  • 1 <= M , Ai<= 1e9

方針・所感

愚直にやるとN個の累積和をとって全てのl,rの組について条件を満たすか判定する、という解法が思い浮かぶがO(N^2)となり、ギリギリ間に合わない。なので、高速化する方法を考えるわけだが、modMで考えて累積和が等しいものがあるならその区間の和はMの倍数となり、modMが等しいものの個数はO(N)で求められる。あとは、その個数xの中からxC2で(l,r)の組を選べば良いので、全体としてはO(N+N)で求められる。

 

細かいところ・デバッグの記録など

  • 累積和を取る時は1-indexedにすると実装が描きやすい
  • 例によってlong longにしなかったせいでxC2がオーバーフローしてしまった、注意
  • 累積和にmintを使いたかったが、mod入力で与えられる場合は別のクラスを定義しなければならないようなので断念

実装

#include <bits/stdc++.h>
using namespace std;
#define rep(i, a, b) for(int i = a; i < b; i++)
typedef long long ll;

int main(){
ll N, M; cin >> N >> M;
vector<ll> A(N+1), SUM(N+1);
rep(i, 1, N + 1) cin >> A[i];
rep(i, 1, N + 1) SUM[i] = SUM[i-1] + A[i];

map<int, ll> MP;
rep(i, 0, N + 1) MP[SUM[i] % M]++;
ll ans = 0;
for(auto e: MP){
ans += e.second * (e.second-1) / 2;
}
cout << ans << endl;
}