AtCoder-ABC194 C - Squared Error【Python解答例】

AtCoder Beginner Contest194のC - Squared ErrorについてPythonの解答例を記事にしていきます。
AtCoder Beginner Contest 194 - AtCoder

AtCoder Beginner Contest194 C - Squared Error

C - Squared Error

問題文概要

長さ N の数列 A が与えられます。各要素同士の差の 2 乗の和 を求めてください。

制約

・2≤N≤3×10^5
・|Ai|≤200
・入力に含まれる値は全て整数

解答例

n = int(input())

A = list(map(int, input().split()))

Sum = [0] * (n + 1)

res = (n-1)*pow(A[0],2)
Sum[0] = A[0]
for i in reversed(range(1,n)):
    Sum[i] = Sum[i+1] + A[i]
    res += (n-1)*pow(A[i],2) -(2 * Sum[i] * A[i-1])

print(res)

解説

問題文通り長さ N の数列 A において各要素同士の差の 2 乗の和を答える問題です。

2重ループにできれば簡単ですが、時間が足りないため工夫する必要があります。

自分の解答は数学的に証明できていませんがACはとれた、という感じなのでご注意ください。

N = 3 , j = 1の場合は(A3^2 - A1^2)^2＋(A2^2 - A1^2)^2となり展開すると2A1^2 + A2^2 + A3^2 -2×A1(A2 + A3)となります。
N = 3 , j = 2の場合は(A3^2 - A2^2)^2となり、A2^2 + A3^2 -2×A2×A3です。
答えとなるその和は2(A1^2 + A2^2 + A3^2) - 2{A1(A2 + A3) + A2×A3}になります。
したがって解答は、それぞれのAの要素の2乗に( N - 1 )をかけた値の合計と、Nからj + 1までの要素の合計に-2×Ajをかけた値の和で求まると予想しました。
（Nからj + 1までの要素の合計は配列Sumとしてループにあわせて更新していきました。）

この方針でACできましたが、証明できていないので参考程度の話になります。

結果オーライではありますが、C問題まで解けてよかったです。
コンテスト中に解けなかったD問題も記事にしていきたいます。

ABC194の関連記事はこちら
ebisuke33.hatenablog.com
ebisuke33.hatenablog.com