問題概要
関数 $f \mathpunct{:} ( \mathbb Z_{ > 0 } )^2 \to \mathbb Z_{ \geq 0 }$ を
\begin{equation*}
f( x, y ) = ( x + y ) \bmod 10^8
\end{equation*}
で定義する.
$n$ 項からなる正整数列 $A = \langle A_1, A_2, \dots, A_n \rangle$ が与えられる.式
\begin{equation*}
\sum_{ i = 1 }^{ n - 1 }\sum_{ j = i + 1 }^{ n } f( A_i, A_j )
\end{equation*}
で定まる値を求めよ.
制約
- $2 \leq n \leq 3 \times 10^5$
- $1 \leq A_i < 10^8$
解法
値を計算するだけなら二重ループを回すことで計算できますが,$\Omega( n^2 )$ 時間かけてしまうと TLE するので工夫する必要があります.経験則的には $O( n \log n )$ 時間ぐらいにしたい気がします(実際,今回はそうします).
問題を難しくしているポイントであり,また一瞬勘違いする*1点として,求めるのは「和の $\bmod$」ではなく「$\bmod$ の和」であるというのがあります.つまり,(雑に?)和を求めてから最後に $\bmod$ をとるのはだめで,$\bmod$ による答えへの寄与をある程度真面目に扱う必要があります.
やや天啓ですが,昔からの格言として「ソートしても答えが変わらないときはとりあえずソートしてみる」というのがあります[要出典].今回の問題は $A$ から二項を選ぶ取り方すべてに渡って値を足し合わせるというものなので, 項の順序は答えに影響せず,$A$ をソートしてもよいです.ソートによって得られる性質を後で使うので,ソート済みの $A$ を $\hat A$ と置いておきます.
二重 $\sum$ の外側の添字 $i$ を固定した状況を考えます.つまり,
\begin{equation*}
\sum_{ j = i + 1 }^{ n } f( \hat A_i, \hat A_j )
\end{equation*}
を考えます.このままだと何も変わってないので何か工夫をしたいわけですが,よくあるパターンとしてはこの状況下では $A_i$ は定数なので $\sum$ の外側に出せると嬉しい気がします.しかし,$f$ の中に入っているのでそのままでは出せません.そこで一旦,やや飛躍しますが
\begin{equation*}
\sum_{ j = i + 1 }^{ n } ( \hat A_i + \hat A_j ) =
( n - i ) \hat A_i + \sum_{ j = i + 1 }^{ n } \hat A_j
\end{equation*}
という値を考えて,本来求めたかった値との差分 $d$ を考えます.$d$ への寄与は $\bmod$ によって発生しているので,$d$ は非負整数 $k$ を用いて
\begin{equation*}
d = k 10^8
\end{equation*}
という形をしているはずです.また,$A_i$ の制約から $\hat A_i + \hat A_j < 2 \times 10^8$ なので,一回の $\bmod$ による変化分は $0$ か $10^8$ のいずれかです.つまり,上記の $k$ とは,$\hat A_i + \hat A_j \geq 10^8$ となる $j$ の個数に等しくなります.
以上をまとめれば,各 $i$ について
\begin{equation*}
( n - i ) \hat A_i + \sum_{ j = i + 1 }^{ n } \hat A_j - k 10^8
\end{equation*}
を求めればよいことになります.ここで,$\sum$ の値は予め $\hat A$ の「累積和」と呼ばれるものを求めておけば,それぞれ $O( 1 )$ 時間で求められるようになります.$k$ の値については,$\hat A$ がソート済みであることから,ある値未満の $j$ では $\bmod$ で和が変化せず,それ以外の $j$ では $10^8$ が引かれるという性質があるので,これを利用します.そのような閾値は $\hat A$ 上での二分探索によって $O( \log n )$ 時間で求めることができます.
従って,各 $i$ について $O( \log n )$ 時間で値を求めることができ,$A$ のソートと合わせても全体で $O( n \log n )$ 時間となり,正答することができます.
コード
constexpr int M = 100'000'000; int main() { IN( int, N ); VLL A( N ); cin >> A; sort( ALL( A ) ); VLL psum( 1 ); partial_sum( ALL( A ), BI( psum ) ); LL res = 0; REP( i, N ) { res += A[i] * ( N - 1 - i ); res += psum[N] - psum[ i + 1 ]; const int pos = lower_bound( begin( A ) + i + 1, end( A ), M - A[i] ) - begin( A ); res -= LL( N - pos ) * M; } cout << res << endl; return 0; }
*1:わたしだけ?