はじめに
ここ3回ほど、AtCoderに参加したばかりの初心者です。お仕事はCを書くことが多いのですが、他の言語にも慣れたくて、Pythonを使って解くことにしています。Pyhton(そしてQiitaも)初心者なので、コード上のアドバイスなどありましたら、コメントいただければ幸いです。
累積和を使わないで解答したのは、累積和というテクニックを知らず、コンテスト中にひねり出した答えだったからです。コンテスト後いくつか解説の記事を見てみたのですが、累積和を使用する解答が多く、記事として残しておくことにしました。
問題概要
ABC177C「Sum of product of pairs」
$N$個の整数$A_{1},....A_{N}$で構成された数列が与えられる。
$1≤i<j≤N$を満たす全ての組$(i,j)(i,j)$についての$Ai×Aj$の和を$mod(10^9+7)$で求めよ。
ABC177Cの解説記事
https://qiita.com/rute_not_route/items/c6fbf3acfcfd93c9258d
考え方
みなさん小学校で習ったので九九の表はご存知だと思いますが、九九の表全体の合計値はご存知ですか？
答えは2045です。シンプルに計算すると、次のように力技で解くことができます。
$$(1\times1) + (1\times2) + \cdots (9\times8)+(9\times9) = 2045 $$
しかし九九の総和は1から9までの合計を計算し、その合計を2乗することで計算することができます。
$$
(1+2+\cdots+9)^2=2045
$$
このアイディアを用いて、次のような方針で解いていきます。
①九九の表のように、入力配列の総和の2乗を計算する。これを$sum$とする。
②九九の表のように並べた場合の対角線要素(例:$A_iA_i$など)の和を計算する。これを$rem$とする。
③答えとして$(sum-rem)/2$を計算する。
イメージとしては、九九表に置ける以下のoで示した部分の合計を計算したいのです。
②でxの部分($rem$)の合計を計算します。
③で全体を2で割ることで、「oエリア」と「-エリア」を区別します。
※「oエリア」と「-エリア」の合計値は同じ。
1 2 3 4 5 6 7 8 9
1 x o o o o o o o o
2 - x o o o o o o o
3 - - x o o o o o o
4 - - - x o o o o o
5 - - - - x o o o o
6 - - - - - x o o o
7 - - - - - - x o o
8 - - - - - - - x o
9 - - - - - - - - x
解答例
上記の考え方を使用しての解答例は、こんな感じです。
n = int(input())
a = list(map(int,input().split()))
max = (10 ** 9) + 7
#入力配列の総和
line = 0
for i in range (n):
line += a[i]
#a11, a22, ...対角線要素の総和
rem = 0
for i in range(n):
rem += a[i]*a[i]
sum = ((line*line) - rem) // 2
ans = sum % max
print(ans)
私がAtCoderに投稿したときには、実行時間は115msでした。