問題概要

整数列 $A$ が与えられる。次の操作を繰り返し行って、 $A$ の最大値と最小値の差が 1 以下となるようにしたい。実現のための操作の最小回数を求めよ。

$i, j$ を選んで、 $A_{i}$ に 1 を足し、 $A_{i}$ から 1 を引く

制約

$1 \le N \le 2 \times 10^{5}$

考えたこと：最終形

「最大値と最小値の差が 1 以下」とは、均等にすることを意味する。たとえば、 $A$ の要素数が $5$ 、総和が $38$ であれば、

$38 = 5 \times 7 + 3$

なので、操作後の配列を $B$ としたとき、順番を無視すれば $B = (7, 7, 8, 8, 8)$ ( $8$ が $3$ 個) となる。基本的には、「最大値と最小値が 1 以下」である状態とは、この並び替えしかない。

$A$ の各要素に対して、 $B$ のどの要素を割り当てるのが最適化かを考える問題となる。

小さいものに小さいものを割り当てる

まず、操作前の $A$ の要素の並びは全く関係ないのでソートして考えることにする。たとえば

$A = (3, 4, 6, 7, 7)$

を考えよう。これは均等にすると $B = (5, 5, 5, 6, 6)$ となる。実は、

小さい $A$ には、小さい $B$ を
大きい $A$ には、大きい $B$ を

割り当てるのが最適となる。つまり、A も B もソートして、その差の和を求めて、2 で割れば OK

コード

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

int main() {
    long long N;
    cin >> N;
    long long S = 0;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        S += A[i];
    }
    sort(A.begin(), A.end());
    
    long long q = S / N, r = S % N;
    vector<long long> B(N);
    for (int i = 0; i < N; ++i) B[i] = q + (i < r);
    sort(B.begin(), B.end());
    
    long long res = 0;
    for (int i = 0; i < N; ++i) res += abs(A[i] - B[i]);
    cout << res/2 << endl;
}