ABC194 C - Squared Error をやった記録

1 min read読了の目安(約1100字

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

https://atcoder.jp/contests/abc194/tasks/abc194_c

問題概要

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

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

解き方

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

  1. まず、-200から200までの数字がそれぞれ何個、数列A_n に含まれているのか(出現回数)をカウントするdictionaryをつくる。
  2. 次に、和を求める。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)  # 整数型にするため // で割っている

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

https://blog.hamayanhamayan.com/entry/2021/03/07/000632
https://note.com/taraba_kanico/n/n8aca58bcd303