📚

ABC194 C - Squared Error をやった記録

2021/03/21に公開

Python

プログラミング

競プロ

tech

本番で解けなかった、入茶したいのでとき直してみた。
Pythonで実装した。公式解説の解法1のほうで実装する。

問題概要

長さNの数列Aの各要素の差の二乗の和を求めよ。

\sum^N_{i=2} \sum^{i-1}_j (A_i -A_j)^2

解き方

差の組み合わせの種類が $400 ^2=1.6 \times 10^5$ 程度であることを利用して解くことができる。

まず、-200から200までの数字がそれぞれ何個、数列 $A_n$ に含まれているのか(出現回数)をカウントするdictionaryをつくる。
次に、和を求める。 $i = -200 \sim 200, j= i+1 \sim 200$ についてそれぞれ　
$(i - j)^2$ $\times$ (iの出現回数) $\times$ (jの出現回数)　
を計算して足し込んでいけば良い

n = int(input())
a = input().split()
dic = {}
for i in range(-200, 201):
    dic[i] = a.count(str(i))
# 和を取る処理
ans = 0
for i in dic.keys():  # このdic.keys() は range(-200,201)で置き換えてもok
    for j in range(i + 1, 201):
        ans += dic[i] * dic[j] * (i - j) ** 2
print(ans)

ちなみに、和を取る処理を $i,j=-200 \ism 200$ とするとTLEで時間切れになり解けなかった。

# 和を取る処理 i,jの範囲を-200~200でとり、最後に2で割る
ans = 0
for i in dic.keys():  # このdic.keys() は range(-200,201)で置き換えてもok
    for j in dic.keys():
        ans += dic[i] * dic[j] * (i - j) ** 2
print(ans // 2)  # 整数型にするため // で割っている

ABC194 C - Squared Error をやった記録

問題概要

解き方

参考にさせていただいた記事

Discussion