ABC313 C - Approximate Equalization 2 (Diff: 681 茶) の解説 [AtCoder][Python]

問題へのリンク https://atcoder.jp/contests/abc313/tasks/abc313_c

問題概要

長さNの整数列Aがある。Aの要素から2つ選び、一方の要素には1を加え、もう一方の要素には1を引く動作を繰り返す。この時、Aの最小値と最大値の差が1になるようにするために必要な最小の操作回数を求めよ。

入力
4
4 7 3 7
出力
3
操作の例

最初 (4,7,3,7)
→ (5,6,3,7)
→ (5,6,4,6)
→ (5,6,5,5)

制約
  • 1 ≦ N ≦ 2×105
  • 1 ≦ A[i] ≦ 109

解説

この問題を理解しやすくするために、直線状に並んだN個の場所にA[i]個、箱が積み上げられているのをイメージする。「目標」は、一番上に積まれた箱を1つ別の場所に移す動作を繰り返し、最も高く積まれた場所と、最も低い場所の個数の差が1個以内にすることである。この時、最も高い場所の箱を1つ取って(1引く)、最も低い場所に1つ積む(1足す)ことを繰り返し、回数を数えていけば良さそうだ。

しかし、個数が多くなってくると、1個1個箱を動かすのは時間がかかる。そのため、各場所で箱を何個動かせば良いのか考える。

操作を繰り返し行い、「目標」を満たした後のことをイメージする。すると、箱は全て同じような高さで、高いほう(最大値)と低いほう(最小値)があり、高い部分は低い部分より1個だけ多く積み上げられているようになっているはずだ。そして、その時積み上げられる高さは全ての箱(Aの総和)の合計をNで割った「平均値」に近い値になっており、端数の部分で高い部分と低い部分が生まれる。例えば、入力例1では、N=4、箱の数(Aの総和)は21個で、平均値は21/4=5.25となるが、高さは整数でないといけないから操作後の高さは(5,6,5,5)のようになり、低い部分(5)と高い部分(6)が生まれている。

操作後に、この低い部分と高い部分が何個生まれるかを考えれば、操作を何回行えば良いのか分かりそうだ。(最小値)+1=(最大値)なので、操作後は先に全て同じ最小値の分積まれ、箱が余った数だけ1個ずつ積まれた状態になる、と考える。入力例1では、操作後のすべての要素は、21/4の切り捨ての5個積まれ、余った1個(21/4の余り)がどこか1か所に積まれる状態になる。これが積まれる場所に関しては、元のAが大きい順に置いていくと、操作回数を最も少なくできる。

そして、全ての場所で(最初の箱の数)-(操作後の箱の数)を計算し、正の数(箱を除くべき回数)をカウントする。

なお、実際は箱の順番、つまりAの要素は入れ替えても影響がない。操作、また最小値・最大値には場所は関係しないからだ。Aを大きい順に並び替えると、先ほどの余った箱を積む場所を大きい順に置いていく部分の実装が楽になる。

実装(Python)

N = int(input())
A = list(map(int, input().split()))

total = sum(A) # Aの総和(箱の個数)
ave = total // N # 総和の平均の切り捨て
L = [ave] * N # 操作後のA[i]の要素(箱の個数)

# 余りの分だけ、大きい順に1を足していく
# (iが小さいほどA[i]が大きくなるようAをソートしている)
for i in range(total % ave):
    L[i] += 1

A.sort(reverse=True) # 要素(箱数)が大きい順に並び変える
ans = 0

for i in range(N):
    if A[i] - L[i] > 0:
        # もし-1する回数が0より大きいなら答えに加える
        # 今回は-1の回数だけを集計している
        ans += A[i] - L[i]

print(ans)