問題概要
長さ $n$ の数列 $A = \langle A_1, A_2, \dots, A_n \rangle$ が与えられる.
以下のクエリが $q$ 個与えられるので,処理せよ:
$i$ 番目のクエリでは,整数 $L_i, R_i$ ($1 \leq L_i \leq R_i \leq n$) が与えられる.以下の値を出力せよ:
\[
\sum_{ l = L_i }^{ R_i } \sum_{ r = l }^{ R_i } \sum_{ j = l }^r A_j
\]
制約
- $1 \leq n \leq 3 \times 10^5$
- $1 \leq q \leq 3 \times 10^5$
- $1 \leq A_i \leq 100$
準備
一般の列 $A$ とその添字 $l, r$ に対して,$A[ l, r ] := \langle A_l, A_{ l + 1 }, \dots, A_r \rangle$ とする記法を採用しています.
解法
一応クエリの意味を確認しておくと,「全ての『与えられた $A$ の連続部分列に含まれる連続部分列についての部分和』の和を求めよ」と言っています.$L_i, R_i$ に対して,$A[ L_i, R_i ]$ に含まれる連続部分列は $\Theta( R_i - L_i + 1 )$ 個あるので,これらにそれぞれ着目することはできません.そこで,各 $A_j$ ($j \in [ L_i, R_i ]$) が何回足されるのかを考えます.
ある $A_j$ ($j \in [ L_i, R_i ]$) が足される回数は,$A_j$ のみを含む $A$ の連続部分列 $\langle A_j \rangle$ から左端点・右端点をそれぞれ伸ばして $A$ の連続部分列 $\langle A_k, A_k + 1, \dots, A_j, \dots, A_{ k' - 1 }, A_{ k' } \rangle$ を作る方法の総数になります.これは左右に伸ばせる長さをそれぞれ数えてかけ合わせればよくて,
\[
( j - L_i + 1 )( R_i - j + 1 )
\]
回となり,
答えへの寄与としては
\[
( j - L_i + 1 )( R_i - j + 1 )A_j
\]
となります.これを $[ L_i, R_i ]$ の全ての $j$ について足し合わせればよいので,$i$ 番目のクエリに対する答えは
\[
\sum_{ j = L_i }^{ R_i } ( j - L_i + 1 )( R_i - j + 1 )A_j
\]
です.これを気合で(過剰に丁寧に)展開して整理すると
\[
\begin{align*}
\sum_{ j = L_i }^{ R_i } ( j - L_i + 1 )( R_i - j + 1 )A_j &=
\sum_{ j = L_i }^{ R_i } ( jR_i - j^2 + j - L_iR_i + jL_i - L_i + R_i - j + 1 )A_j \\
&= \sum_{ j = L_i }^{ R_i } ( -j^2 + jL_i + jR_i + R_i - L_i- L_iR_i + 1 )A_j \\
&= \sum_{ j = L_i }^{ R_i } ( -j^2 + j( L_i + R_i ) + ( R_i - L_i- L_iR_i + 1 ) )A_j \\
&= \sum_{ j = L_i }^{ R_i } ( -j^2A_j + ( L_i + R_i )jA_j + ( R_i - L_i- L_iR_i + 1 )A_j ) \\
&= \sum_{ j = L_i }^{ R_i } -j^2A_j +
\sum_{ j = L_i }^{ R_i }( L_i + R_i )jA_j +
\sum_{ j = L_i }^{ R_i }( R_i - L_i- L_iR_i + 1 )A_j \\
&= -\sum_{ j = L_i }^{ R_i }j^2A_j +
( L_i + R_i )\sum_{ j = L_i }^{ R_i }jA_j +
( R_i - L_i- L_iR_i + 1 )\sum_{ j = L_i }^{ R_i }A_j
\end{align*}
\]
となります.ここで,各 $\sum$ の係数は $L_i, R_i$ が与えられれば定まってかつ $O( 1 )$ 時間で計算できるので,各 $\sum$ の値を高速に計算できれば問題を解けます.ちょっと面食らいますがよく考えると,$\langle 1^2A_1, 2^2A_1, \dots, n^2A_n \rangle$ と $\langle 1A_1, 2A_1, \dots, nA_n \rangle$ の $2$ の列を新たに考えて,(元の $A$ も含めて)累積和をそれぞれ取っておけば,$\langle \Theta( n ), O( 1 ) \rangle$ 時間で処理できます.
上記の方法の計算量は前処理に $\Theta( n )$ 時間,各クエリに $O( 1 )$ 時間かかります.従って全体で $\Theta( n + q )$ 時間であり,TL に間に合います.
コード
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, Q ); VLL A( N ); cin >> A; VT< Psum< LL > > psums; { VVLL B( 3, VLL( N ) ); REP( i, N ) { B[0][i] = A[i]; B[1][i] = ( i + 1 ) * A[i]; B[2][i] = LL( i + 1 ) * ( i + 1 ) * A[i]; } for_each( ALL( B ), [&]( const auto &b ){ psums.EB( b ); } ); } REP( Q ) { IN( LL, L, R ); cout << -psums[2].sum( L - 1, R ) + ( L + R ) * psums[1].sum( L - 1, R ) + ( R - L - LL( L ) * R + 1 ) * psums[0].sum( L - 1, R ) << '\n'; } cout << flush; return 0; }