E - Mod Sigma Problem Editorial by en_translator
The typical approach to represent subarray sums employs cumulative sums. Let \(S_i = (A_1+A_2+\dots+A_i) \mathbin{\mathrm{mod}} M\) be the cumulative sums (modulo \(M\)) of the sequence \(A\); then the sought value can be represented as:
\[ \sum_{1 \leq l \leq r \leq N} \left( \left(\sum_{l \leq i \leq r} A_i\right) \mathbin{\mathrm{mod}} M \right) = \sum_{1 \leq l \leq r \leq N} (S_r - S_{l-1}) \mathbin{\mathrm{mod}} M, \]
where we define \(S_0\coloneqq 0\).
Since \(0 \leq S_{l-1},S_r < M\), we can write as
\[ (S_r - S_{l-1}) \mathbin{\mathrm{mod}} M = S_r - S_{l-1} + \begin{cases} 0 & (S_{l-1} \leq S_r) \\ M & (S_{l-1} > S_r),\end{cases} \]
which no longer involves \(\mathrm{mod}\).
Let \(X_r = \) (the number of \(l=1,2,\dots,r\) with \(S_{l-1} > S_r\)) and apply the equation above to rewrite the sought value as:
\[ \sum_{r=1}^N \sum_{l=1}^r (S_r - S_{l-1}) \mathbin{\mathrm{mod}} M = \sum_{r=1}^N \left( S_r \times r - \sum_{l=1}^r S_{l-1} + M \times X_r \right). \]
Then \(\sum_{l=1}^r S_{l-1}\) is a cumulative sum of \(S\), which can be evaluated easily. All that left is to find \(X_r\) for \(r=1,2,\dots,N\) fast enough for each \(r=1,2,\dots,N\).
\(X_r\) can be computed using a Fenwick tree. A Fenwick maintains a sequence \(B=(B_1,B_2,\dots,B_M)\) and supports the following operations in \(O(\log M)\) time each:
- Given \(i\) and \(x\), set \(B_i \leftarrow B_i+x\).
- Given \(l\) and \(r\), find \(B_l + B_{l+1} + \dots + B_r\).
Here, if we define \(B_x\) as the number of \(l\) with \(S_{l-1}=x\), we can compute \(X_r\) by the following algorithm:
- Initialize with \(B_x = 0\).
- For each \(r=0,1,\dots,N\), do the following:
- Compute \(X_r = B_{S_r + 1} + B_{S_r + 2} + \dots B_{M-1}\).
- Add \(1\) to \(B_{S_r}\).
It is convenient to use Fenwick tree provided in AtCoder Library.
The overall complexity is \(O(N\log M + M)\).
Sample code (Python)
from itertools import accumulate
from atcoder.fenwicktree import FenwickTree
N, M = map(int, input().split())
A = list(map(int, input().split()))
S = [0] + [x % M for x in accumulate(A)]
ft = FenwickTree(M)
s = 0
ans = 0
for r in range(N + 1):
ans += S[r] * r - s + ft.sum(S[r] + 1, M) * M
s += S[r]
ft.add(S[r], 1)
print(ans)
posted:
last update: