ABC 177 C - Sum of product of pairs

2020年08月30日

AtCoder Beginner ContestのC問題だけやるシリーズをちょっと始めてみようかな〜と思います。なんでC問題だけ解くかというと、初学者がぎりぎり解けそうで解けないのがC問題くらいだから。

使う言語はrubyにします。みんながよく使ってるjavascriptかrubyで迷ったんですけどjavascriptは入力処理がマゾすぎるのでやめました。ちゃんと競プロやりたい人はC++がオススメですけど。

今日やるのは ABC 177 C - Sum of product of pairs

問題みて最初に思うこと

1 <= i < j <= N を満たす全ての組 (i, j) について

ペアは O(N ** 2) 個あることがわかるけど制約には

2 <= N <= 2 * 10 ** 5

って書いてあるので二重ループすると計算量が O(4 * 10 ** 10) になっちゃってTLEしますね。ここを上手になんとかできるかどうかがこの問題のポイントになってるみたい。

Ai * Aj の和を mod 10 ** 9 + 7 で求めてください

これは大きい整数がでるときにあるあるなやつで、オーバーフローでWAになるのイヤなのでちょこちょこmodするだけ。

解き方を考える

modのことはしばらく忘れて話をするんですけど。数学的にちゃんとしたのじゃなくて、解くときの気持ちに近いように書きます。

ans = (Ai * Aj) の総和（と一時的に呼ぶ）
Aの要素がansの値に影響するのは、どれかのペアのAiまたはAjであるとき
具体的に10番目の要素A10をみて、これがAjとして登場するペアの和は
(A1 * A10) + (A2 * A10) + … + (A9 * A10)
= (A1 + A2 + … + A9) * A10
配列を左から見てって、そこまでの総和を覚えてると O(1) で計算できる
ans = (j=1のとき) + (j=2のとき) + … + (j=Nのとき)
全体で O(N) でとけるのでTLEにならない

みたいなことを考えればオッケーです。

解答コード

MOD = 10 ** 9 + 7

N = gets.to_i
A = gets.split(" ").map(&:to_i)

sum = 0
product_sum = 0

A.each do |a|
  product_sum += sum * a
  product_sum %= MOD
  sum += a
  sum %= MOD
end

puts product_sum

それぞれの要素について、それより左にある要素の総和と掛け算したのを足してく〜というのをそのまま書きました。ループがひとつだけでオッケーなところが気持ちよい。

今回のC問題は文法構造がややこしいところはなくて、考察がちょっといる感じのでした。アルゴリズムとしては『累積和』と呼ばれてるやつで、めっちゃよく出てくるやつなので、覚えておくとよいことあります。