問題概要
$n$ 項からなる非負整数列 $A = \langle A_1, A_2, \dots, A_n \rangle$ と正整数 $m $ が与えられる．
$\mathcal I = \{ [ l, r ) \mid l, r \in \{ 1, 2, \dots, n + 1 \}, l < r \}$ として，次を求めよ：
\begin{equation*}
\sum\limits_{ I \in \mathcal I } \left( \left( \sum\limits_{ i \in I } A_i \right) \bmod m \right)
\end{equation*}
制約
- $1 \leq n \leq 2 \times 10^5$
- $1 \leq m \leq 2 \times 10^5$
- $0 \leq A_i \leq 10^9$
解法
数列 $A$ と整数 $l, r \in \{ 1, 2, \dots, |A| + 1 \}$ ($l < r$) に対し，数列 $\langle A_l, A_{ l + 1 }, \dots, A_{ r - 1 } \rangle$ を $A[ l, r )$ と書くことにします．また，数列 $A$ に対し，その総和を $\sum A$ と略記します．
$\mathcal I$ の元が $\Theta( n^2 )$ 個あるので，それぞれの区間和を愚直に計算することはできません．$\mathord{ \bmod }$ をとっていなければ各 $A_i$ が答えに何回寄与するかを考える（$l \leq i, i < r$ になるので $l, r$ の取り方の総数を算数できる）ことで求められますが，今回の問題を解くには，$\mathord{ \bmod }$ をとる処理が答えに与える影響を考える必要があり，各 $A_i$ に対してその回数を高速に計算するのは（各 $l, r$ に対して部分和を計算する必要がありそうで）むずかしそうに見えます．
やや知識の話ですが，数列の区間和を求める方法として「累積和」と呼ばれるものが知られています．これは，数列 $A$ に対して数列 $P$ を
\begin{equation*}
P_i = \begin{cases}
0 & \text{($i = 0$)} \\
P_{ i - 1 } + A_i & \text{(otherwise)}
\end{cases}
\end{equation*}
を定義し利用するものです．各 $r \in \{ 1, 2, \dots, n \}$ に対して $\sum A[ 0, r ) = P_r$ であることから，$\sum A[ l, r ) = \sum[ 0, r ) - \sum[ 0, l )$ と分解して，$P_r - P_{ l - 1 }$ で区間和を $\langle \Theta( |A| ), O( 1 ) \rangle$*1で求められます．
累積和を踏まえて問題で与えられている式を観察します．区間和を上記のように分解して，後から $\mathord{ \bmod }$ を取っていることからその前にそれぞれ $\mathord{ \bmod }$ をとってもよいということを踏まえると，外側の $\mathord{ \sum }$ の内側は
\begin{align*}
\sum\limits_{ i \in [ l, r ) } A_i \bmod m
&= \sum A[ l, r ) \bmod m \\
&= ( \sum A[ 0, r ) - \sum A[ 0, l ) ) \bmod m \\
&= ( \sum A[ 0, r ) \bmod m - \sum A[ 0, l ) \bmod m ) \bmod m
\end{align*}
と変形できます．$\mathord{ \bmod }$ の定義を思い出すと整数 $x$ と正整数 $m $ に対して $x \bmod m $ は $x$ を整数 $q, r$ ($0 \leq r < m $) を用いて $x = q m + r$ と表したときの $r$ を求める操作です．上の式で最後に $\mathord{ \bmod }$ をとるときのことを考えると，括弧の内側の減算のオペランドがともに $m $ 未満なので，$-m < x < m $ であるような $x$ に対して $x \bmod m $ を計算していて，この場合，
\begin{equation*}
x \bmod m = \begin{cases}
x & \text{($x \geq 0$)} \\
x + m & \text{(otherwise)}
\end{cases}
\end{equation*}
となります．
以上を踏まえた上で，問題で与えられた式の外側の $\mathord{ \sum }$ （区間和をとる範囲を動かしている部分）において区間の右端点を固定した場合を考えます．この場合，式は
\begin{equation*}
\sum\limits_{ l = 0 }^{ r - 1 } \left( \left( \sum A[ l, r ) \right) \bmod m \right)
= \sum\limits_{ l = 0 }^{ r - 1 } \left( \left( \sum A[ 0, r ) \bmod m - \sum A[ 0, l ) \bmod m \right) \bmod m \right)
\end{equation*}
となります．外側の $\mathord{ \bmod }$ を上述のように書き換えると，カッコ内が負になるような $l$ の個数を $k$ として
\begin{equation*}
\sum\limits_{ l = 0 }^{ r - 1 } \left( \sum A[ 0, r ) \bmod m - \sum A[ 0, l ) \bmod m \right) + k m
\end{equation*}
となります．更に $\sum A[ 0, r )$ の部分は $r$ 回固定値を足しているので $\mathord{ \sum }$ の外に出して，
\begin{equation*}
r \sum ( A[ 0, r ) \bmod m ) - \sum\limits_{ l = 0 }^{ r - 1 } \left( \sum A[ 0, l ) \bmod m \right) + k m
\end{equation*}
になります．ここで
\begin{equation*}
\sum\limits_{ l = 0 }^{ r - 1 } \left( \sum A[ 0, l ) \bmod m \right)
\end{equation*}
の部分は $i$ 項目が $A[ 0, i ) \bmod m $ であるような列の累積和となっているので，$r$ を昇順に動かすとすれば差分を $O( 1 )$ 時間で計算できます．残っているのは $k$ ですが，$k$ は $A[ 0, r ) \bmod m $ より大きい $A[ 0, l ) \bmod m $ の個数とも言い換えられるので，区間 $[ 0, m )$ に属する各整数について，その出現回数を覚える列を作って区間和を求めることで計算できます．この列の更新と区間和を高速に処理できるデータ構造としては Fenwick Tree があり，列上の一点の更新と区間和をそれぞれ $O( \log m )$ 時間で処理できます．
まとめると，$r$ を昇順に動かしながら累積和と Fenwick Tree を随時更新することで各 $r$ による答えへの寄与の計算に必要な値を揃えることができて，その計算量は全体で $O( n \log m )$ 時間となります．
コード
// Fenwick Tree ( Binary Indexed Tree ) template < typename T > class FenwickTree; // 中身省略 // FenwickTree( array size ) // int sum( 0-indexed pos ) // int sum( [ l, r ) // int add( pos, value ) // int change( pos, value ) // int lower_bound( value ) template < typename T > struct Psum { std::vector< T > psum; Psum( const std::vector< T > &A ) : psum( 1 ) { std::partial_sum( std::begin( A ), std::end( A ), std::back_inserter( psum ) ); return; } T operator[]( const int i ) const { return psum[ i + 1 ]; } T sum( const int l, const int r ) const { return psum[r] - psum[l]; } }; int main() { IN( int, N, M ); VLL A( N ); cin >> A; LL res = 0, s = 0; FenwickTree< int > fenwick( M ); const Psum psum( A ); REP( r, N ) { res += LL( r + 1 ) * ( psum[r] % M ) - s + LL( M ) * fenwick.sum( psum[r] % M + 1, M ); s += psum[r] % M; fenwick.add( psum[r] % M, 1 ); } cout << res << endl; return 0; }
