Atcoder Beginner Contests 194 C - Squared Error (Python 3.8.2)

皆さんお久しぶりです！！
今日はなんとなく記事を書きたい気分になったので書きたいと思います。

問題はこちら！！

atcoder.jp

考察

見た瞬間に「これ類題あるやん〜〜〜」ってなりましたね。

その類題はこちら

atcoder.jp

これを解いた際には少し時間がかかってしまいましたが、結論から言うとこの問題を解くカギは累積和です！

解法

とりあえず眺めててもわからないので具体例を考えてみる。

n = 4のとき

$\displaystyle \begin{eqnarray*} \sum_{i = 2}^{4} \sum_{j = 1}^{i - 1} (A_i - A_j) ^ 2 &=& (A_1 - A_2)^2 + (A_1 - A_3)^2 + (A_1 - A_4)^2 + (A_2 - A_3)^2 + (A_2 - A_4)^2 + (A_3 - A_4)^2 \\ & = & 3{A_1}^2 + 3{A_2}^2 + 3{A_3}^2 + 3{A_4}^2 -2(A_1 A_2+A_1 A_3+A_1 A_4) - 2(A_2 A_3 + A_2A_4) - 2(A_3A_4) \\ & = & 3({A_1}^2 + {A_2}^2 + {A_3}^2 + {A_4}^2) -2A_1(A_2 + A_3 + A_4) -2A_2(A_3 + A_4) - 2A_3(A_4) \end{eqnarray*}$

といったような式が出てきます。
前半の２乗の部分はO(N)で計算することができ、後ろの一つずつ項が小さくなっていくものは前処理で累積和を取っておくとO(N)で計算することが可能です。

一般にnのときでやってみると以下のようになります。

$\displaystyle \begin{eqnarray*} \sum_{i = 2}^{N} \sum_{j = 1}^{i - 1} (A_i - A_j) ^ 2 &=& \left\{(A_1 - A_2)^2 + (A_1 - A_3)^2 +...+ (A_1 - A_N)^2 \right\} + \left\{(A_2 - A_3)^2 + (A_2 - A_4)^2 +...+ (A_2 - A_N)^2\right\} +...+ \left\{(A_{N-1} - A_N)^2\right\} \\ & = & (N-1){A_1}^2 + (N-1){A_2}^2 +...+(N - 1){A_N}^2 \\ & = & (N - 1)({A_1}^2 + {A_2}^2 + {A_3}^2 +...+{A_{N - 1}^2} + {A_N}^2) -2A_1(A_2 + A_3 +...+ A_N) -2A_2(A_3 +...+ A_N) -...- 2A_{N - 1}(A_N) \end{eqnarray*}$

が得られるのであとはこれを累積和を用いて計算するだけです！！

提出コード

n = int(input())
a = list(map(int,input().split()))
b = [0] * n

# Bnを後ろから累積和を取る
for i in range(n):
    b[i] = a[i]
b = b[::-1]
wa = [0] * n
wa[0] = b[0]
for i in range(n - 1):
    wa[i + 1] += wa[i] + b[i + 1]

x = y = z = 0
for i in range(n - 1):
    x += (n - 1 - i) * (a[i] ** 2)

for i in range(n - 1):
    y -= 2 * a[i] * (wa[-2 - i])

for i in range(n  - 1):
    z += (i + 1) * (a[i + 1] ** 2)

print(x + y + z)

最後に

今回の問題は前回の復習が生きたおかげで速く解くことができたのでよかったです！！
今後も数問にはたくさん挑戦して慣れていきたいものです。

間違いなどがあれば指摘の方をよろしくお願いします。
ありがとうございました！