ABC156-C Rally 解説
\(O(N)\) で解きます。
問題
リンクはこちら
概要
$$ \underset{P}{\min} \sum _ {i=1} ^ {N} (X _ i - P) ^ 2 $$ を求めなさい。
解法
制約が \( N \leq 100,\ X _ i \leq 100\) と小さいので、 \(X\) の取りうる範囲すべてを調べても \(O(NA) \, (A = \max ( X ) - \min ( X ))\) で解くことができます。しかし、少し数学的な考察を挟むと、最小値を取る \(P\) の候補を高々2つに絞ることができます。
$$ f(P) = \sum _ {i=1} ^ {N} (X _ i - P) ^ 2 = \sum _ {i=1} ^ {N} (X ^ 2 _ i - 2 X _ i P + P ^ 2) $$
とおくと、
$$ \begin{aligned} f ' (P) &= \sum _ {i=1} ^ {N} (X ^ 2 _ i - 2 X _ i P + P ^ 2) ' \\ &= \sum _ {i=1} ^ {N} ( 0 - 2 X _ i + 2 P ) \\ &= 2 ( NP - \sum _ {i=1} ^ {N} X _ i ) \end{aligned} $$
となります。
よって、 \( S = \sum X _ i\) とすると \(f(P)\) の増減表（懐かしい響き！）は以下のようになります。
$$ \begin{array}{|c| ccc|}\hline x & \cdots & \frac{S}{N} & \cdots \\ \hline f'(P)& - & 0 & + \\ \hline f(P) & \searrow & f(\frac{S}{N}) & \nearrow \\ \hline \end{array} $$
つまり、 \(f(P)\) は \(X\) の平均値 \(\overline{X} = \frac{S}{N} \) で最小値をとることがわかります。ただし、この問題の \(P\) は整数の値しか取りえないので、 \( f ( \lfloor \overline{X} \rfloor ) , f ( \lceil \overline{X} \rceil ) \) のうち小さいほうを出力します。
以上をまとめるとこのようなコードになります。
N = int(input()) X = list(map(int, input().split())) def cost(p): res = 0 for x in X: res += (x - p)**2 return res m = sum(X) // N print(min(cost(m), cost(m+1)))
余談ですが、 \(P\) に \(\overline{X}\) が入ることを念頭において冒頭の式を見ると、この問題は \(X\) のデータに対して最小二乗法を用いて最確値を求めることに相当している、と気づけます。統計学の素養のある人なら問題を見た瞬間に「あ、残差二乗和の最小値だから \(P = \overline{X}\) だ」と察知できたりするのでしょうか。
