ABC194のC問題を解く

残念ながら解けなかったABC194のC問題を反省しつつ、理解しながらふりかえります。
数式展開も丁寧に説明しました。
問題は長さNの数列が与えられて、各要素どうしの2乗の和を求める問題です。

atcoder.jp

Nが最大で 3 * 10^5なので、これは単純に組み合わせで計算するとアウトになるので、何かしら工夫しないといけません。
私が悪かったのは、 A_i \leq |200| の制約に着目できなかったことと、数学的解法でいくなら、その路線に舵をきりきれずに中途半端に数学的解法を諦めてしまったことです。(数式手で書いてて馬鹿なミスしたのも痛かった。。。)

公式解説の解法1

 A_i \leq |200|の制約に着目して、それぞれの数字の出現した回数を数えて、
2つの数字の組み合わせの絶対値とそれぞれの出現回数を掛けて、合計します。

例えば、数列が[1, 5, 5, 5, 3, 3]だとすると、5が3回出現して、3が2回出現してます。
全通りの組み合わせの中で、5 - 3の組み合わせが出現する回数は、5が出現する回数と3が出現する回数をかけた数になります。
なので、5と3の組み合わせについての差の2乗の和は | 5 - 3 | * 3 * 2 = 12となります。

出現回数を記録する配列は A_i \leq |200|ということで A_iは負の値もとるので、200を足してあげて0〜400の添字に記録します。

count = [0 for _ in range(401)]

for a in a_list:
    count[a + 200] += 1

全通りの組み合わせの差の2乗の和をもとめます。

N = int(input())
a_list = list(map(int, input().split()))

count = [0 for _ in range(401)]

for a in a_list:
    count[a + 200] += 1

ans = 0

for i in range(401):
    for j in range(i, 401):
        ans += count[i] * count[j] * (abs((i - 200) - (j - 200)))**2

print(ans)

数が小さいところに着目するのは鉄則なのに、何をやってたんだ。。。

公式解説の解法2

数学的に導いて解く方法です。

まず最初の式ですが、 (A_i - A_i)^2 = 0なのと、全通りの差分の和を2で割った値が答ということに着目して
 \frac{1}{2} \sum_{i=1}^{N} \sum_{j=1}^{N} (A_i - A_j)^2となります。
数列が[0, 1, 2, 3, 4]の場合は以下のようなイメージです。

全体の和を2でわる理由の図
全体の和を2でわる理由

そこから2乗の部分を展開します。
 \frac{1}{2} \sum_{i=1}^{N} \sum_{j=1}^{N} (A_i - A_j)^2 \\ \Leftrightarrow \frac{1}{2} \sum_{i=1}^{N} \sum_{j=1}^{N} (A_i^2 + A_j^2 + 2A_iA_j)

あと \sum_{i=1}^{N}A_i^2 \sum_{j=1}^{N}A_j^2は同じ数列なので等しい値になります。
なので、  \sum_{i=1}^{N} \sum_{j=1}^{N} (A_i^2 + A_j^2)の部分は
 \sum_{i=1}^{N} \sum_{j=1}^{N} (2A_i^2)と置き換えることができます。
添字のjに全然関係ない数字がN個あるつまり \sum_{j=1}^{N} a = N*aということがわかると、

 \sum_{i=1}^{N} \sum_{j=1}^{N} (2A_i^2) \\ \Leftrightarrow 2N \sum_{i=1}^{N} A_i^2

となります。

一度整理すると、
 \frac{1}{2} (2N\sum_{i=1}^{N} A_i^2 - 2 \sum_{i=1}^{N} ( A_i \sum_{j=1}^{N} A_j))になります。

次に最後の項に着目します。
少し分解して表現すると
 2 \sum_{i=1}^{N} ( A_i \sum_{j=1}^{N} A_j) \\ \Leftrightarrow 2 \sum_{i=1}^{N} ( A_i (A_1 + \cdots + A_N)) \\ \Leftrightarrow 2 (\sum_{i=1}^{N} A_i)^2

もう少しわかりやすいイメージ

最後の項の数式展開
最後の項の数式展開

まとめると以下のようになります。

 \frac{1}{2} (2N\sum_{i=1}^{N} A_i^2 - 2 \sum_{i=1}^{N} ( A_i \sum_{j=1}^{N} A_j)) \\ \Leftrightarrow \frac{1}{2} (2N\sum_{i=1}^{N} A_i^2 - 2(\sum_{i=1}^{N} A_i)^2) \\ \Leftrightarrow N\sum_{i=1}^{N} A_i^2 - (\sum_{i=1}^{N} A_i)^2

大学の統計の授業で見たことある展開方法でした。先生、ごめんなさい。

コードに落とし込むと以下のようになります。

N = int(input())
a_list = list(map(int, input().split()))

double_sum = 0
for a in a_list:
    double_sum += a**2

ans = N * double_sum - (sum(a_list)**2)

print(ans)

レーティングがかなり落ちてしまい、ショックなのと、悔しいです。
自分の実力のなさを認めながら向き合いつつ、鍛錬していこうと思います。