AtCoder ABC105: D - Candy Distribution

区間和は「累積和の差分」で考えてみましょう。水色diff です。

問題

数列 $A_1, \cdots, A_{N}$ の区間 $(l, r)$ で $A_l + A_{l+1} + \cdots + A_r$ が $M$ の倍数であるものの数を求めよ。

制約

$1 \le N \le 10^5$
$1 \le M \le 10^5$
$1 \le A_i \le 10^9$

考察

$\{A_i\}$ の累積和 $\{S_i\}$ を計算し、

$\begin{aligned} A_l + A_{l+1} + \cdots + A_r = \text{(}M\text{の倍数)} &\Longleftrightarrow S_{r} - S_{l-1} = \text{(}M\text{の倍数)} \\ &\Longleftrightarrow S_{r} \equiv S_{l-1} \quad \text{(mod} M \text{)} \end{aligned}$

とします。

求める値は、 $\{S_i\}$ の剰余の数列 $\{T_i\}$ で、同じ値が何度出てくるかの総和と等しくなります。

Pythonでは collections.Counter が便利です。 values() で出現回数が取得できます。

実装例

N, M = map(int, input().split())
A = list(map(int, input().split()))

S = [0] + list(itertools.accumulate(A))
T = [s % M for s in S]
counter = collections.Counter(T)
ans = sum(n * (n-1) // 2 for n in counter.values())
print(ans)