AtCoder Beginner Contest 177 C – Sum of product of pairs を解いた記録

問題の概要

\(N\)個の整数\(A_1, A_2, \dots, A_N\)が与えられる.
\(1 \leq i < j \leq N\)を満たす全ての組\((i, j)\)について\(A_i \times A_j\)を求め,その和を\(10^{9}+7\)で割った余りを出力せよ.

問題へのリンク

制約

  • \(2 \leq N \leq 2 \times 10^5\)
  • \(0 \leq A_i \leq 10^9\)
  • 入力は全て整数

解法

愚直に全ての組\((i, j)\)について\(A_i \times A_j\)を求めていると,明らかに間に合いません.

さて,\(1 \leq i < j \leq N\)という制約より

  • \(i=1\)とすると,対応する\(j\)は\(j=2, 3, 4, \dots, N\)
  • \(i=2\)とすると,対応する\(j\)は\(j=3, 4, \dots, N\)
  • \(i=3\)とすると,対応する\(j\)は\(j=4, \dots, N\)
  • (中略)
  • \(i=N-2\)とすると,対応する\(j\)は\(j=N-1, N\)
  • \(i=N-1\)とすると,対応する\(j\)は\(j=N\)

となります.\(i\)を適当に固定して考えると,次式が成り立ちます.
\[
A_i \times A_{i+1} + A_i \times A_{i+2} + \dots A_i \times A_N = A_i \times (A_{i+1} + A_{i+2} + \dots + A_{N})
\]

\(S_i = A_i \times (A_{i+1} + A_{i+2} + \dots + A_{N})\)と置くと,求める答えは\(S_1 + S_2 + \dots + S_{N-1}\)を\(10^9+7\)で割った余りです.
各\(i\)について\(S_i\)を\(O(1)\)で計算できれば,全体の計算は\(O(N)\)で行えるので制限時間内に答えを計算できます.

ここでは\(S_i\)を\(O(1)\)で計算できるようにするために「累積和」というテクニックを使います.次のような配列\(sum\)を用意します.
\[
sum[i] = (A_i + A_{i+1} + \dots A_N) \bmod 10^9+7
\]
\(2 \leq i \leq N\)となる各\(i\)について\(sum[i]\)を計算する作業は\(O(N)\)で行えます.よって,前処理\(O(N)\),全体の計算\(O(N)\)で答えを計算できます.

実装例

感想

ABC177では前回に引き続き5完できましたが,全人類がDとEを高速でACしていたので緑パフォで終わり,レートも微妙に溶けました.辛いです.