はじめに

こんにちは。今回はABC043のC問題の解説をやっていこうと思います。
問題文 https://atcoder.jp/contests/abc043/tasks/arc059_a

解説

公式解説では揃える整数を-100から100で全探索して最小コストを調べる、という解法でした。計算量は $O(N)$ ですが、-100から100で全探索しているので定数倍が重いです。
僕はこの問題を見た時に、ノートに数式を書きました。それをもとに解説していきます。
まず、最終的に揃える整数を $x$ とします。その時のコストの総和は、

\sum_{i=1}^N (a_i - x)^2

と表せます。これを展開すると、

\sum_{i=1}^N a_i^2 - 2a_ix + x^2

=N*x^2 - 2x*\sum_{i=1}^N a_i + \sum_{i=1}^N a_i^2

となります。これは $x$ の二次関数となっています(以降 $f(x)$ と呼びます)。そして、 $1 \le N$なので $f(x)$ のグラフは下に凸です。よって、 $f(x)の極値 = f(x)の最小値$となり、 $f(x)$ の極値を求める問題になりました。(正確には、揃える数が整数の範囲しか動けないので、$f(x)の最小値(xは実数)=答え$とは限りません)
極値は $f'(x)$ ($f(x)$ の導関数)が $0$ となるような $x$ を $f(x)$ に代入した数です。さっそく微分していきます。

f'(x)=2Nx - 2 * \sum_{i=1}^N a_i

となりました。これが $0$ になれば良いので、

2Nx - 2 * \sum_{i=1}^N a_i = 0

2Nx = 2 * \sum_{i=1}^N a_i

x = \frac{1}{N}*\sum_{i=1}^N a_i

となり、これが$ f(x) $が最小になる $x$ です。ですが、先程述べた通り、 $f(x)の最小値(xは実数)=f(x)の最小値(xは整数)$ となるとは限らないので、 $x$ が実数の時の最適な $x$ から整数の時の最適な $x$ にする方法を考えましょう。
まず、実数の時の最適な $x$ が整数である時、整数の時も変わりません。
実数の時の最適な $x$ が有理数(自然数の和/自然数は有理数となるので)の時はどうすればよいでしょうか。結論から言うと、実数の時の最適な $x$ を小数表記した数を小数第一位で四捨五入した値が整数の時の最適な $x$ となります。詳しい証明はしません。
提出したコードはこちらです(ライブラリ無し)

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

int cost(int a, int b) {
  int sa = a - b;
  return sa * sa;
}

int main() {
  int N;
  cin >> N;
  vector<int> a(N);
  ll sum = 0;
  for(int &o : a) {
    cin >> o;
    sum += o;
  }
  double x = (double)sum / N;
  ll cost_sum = 0;
  for(int o : a) {
    cost_sum += cost(o, sum / N + (x * 2 > sum / N * 2 + 1));
  }
  cout << cost_sum << endl;
}

22行目の $2x > 2*\frac{sum}{N}+1$ って何だ、と思う方へ
これは「整数の時の最適な $x$ は $\lfloor x \rfloor + 1$ か」の判定の式を移項してわかりやすくしたものです。(同じ意味ではないですが簡単にすると「$x$ を四捨五入した数を $X$ とした時、$x < X$ か?」の判定です)

計算量は公式解説と同じく $O(N)$ ですが、だいぶ定数倍が軽くなりました。($200$ 倍くらいから $3$ 倍くらいになった)
ABC043は試験管無し(知らない人はAtCoder ProblemsのABC001やARC001を見ましょう)になって2回目のコンテストなので、問題傾向は今と比べてかなり変わっていますが、数学を使って高速に解ける問題をまた解いてみたいです。

おまけ

今回の問題の答えを出すものを関数グラフで作りました。
https://www.desmos.com/calculator/btsp3oo1sq?lang=ja
使い方

Inという配列に入力を1つずつ入れてください。$N$ は入れなくていいです。
自動的に計算されて、グラフには $f(x)$ が、Outという数字には答え(最小コスト)が出力されます。