203AC - ABC043C - Be Together - back-of-the-envelope

ABC043-C (200 points)

問題

$N$ 個の整数 $a_1, \ldots, a_N$ に対して, $\sum (x - a_i)^2$ を最小化する整数 $x$ を求める.

$1 \leq N \leq 100$
$-100 \leq a_i \leq 100$
入力は全て整数

方針

目的関数 $f(x) = \sum (x - a_i)^2$ を微分すると $f'(x) = 2 (Nx - \sum a_i)$ となるので, $f(x)$ は下に凸な関数であり, $x = \frac{1}{N} \sum a_i$ で最小値を取る. $x$ は整数なので, この近くの整数だけ調べれば良い.

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

ans1 = sum(A) // n + (sum(A) % n != 0)
ans2 = ans1 - 1

cost1 = 0
cost2 = 0

for i in range(n):
    cost1 += (A[i] - ans1) ** 2
    cost2 += (A[i] - ans2) ** 2

print(min(cost1, cost2))