【AtCoder】ABC313 – C問題について解説【Python】

そまちょブログのそまちょ（@somachob）です。

この記事は、AtCoder Beginner Contest 313 の C問題についての解説です。

B問題について解説した記事もあります。

C – Approximate Equalization 2

問題へのリンク

入力例１をもとに解説します。

4
4 7 3 7

この整数列についていずれかの数字を１減らし、いずれかの数字を１増やす操作をします。

最終的に最小値と最大値が１以下にするための最小操作回数を求めます。

たとえば、次の３回の操作で最小値と最大値の差を１以下にすることができます。

• i=2, j=3 として操作すると（4, 6, 4, 7）になる。
• i=4, j=1 として操作すると（5, 6, 4, 6）になる。
• i=4, j=3 として操作すると（5, 6, 5, 5）になる。

規則を見つける

この問題の操作では１減らして、１増やすので、「操作前の整数列の合計」と「操作後の整数列の合計」は変わりません。

何度操作をしても整数列の合計は同じです。

• （4, 7, 3, 7）　⇒　4+7+3+7 ＝ 21
• （4, 6, 4, 7）　⇒　4+6+4+7 = 21
• （5, 6, 4, 6）　⇒　5+6+4+6 = 21
• （5, 6, 5, 5）　⇒　5+6+5+5 = 21

このように A_i を１減らし、A_j を１増やすという操作を何度しても、整数列の合計が変わっていないことが確認できます。

このことから最終的な整数列については、A の合計から求めることができます。

最終的な整数列

入力例１の４つの数字の整数列 A ＝（4, 7, 3, 7）の合計は 21 です。

４つの数字で合計が 21 。そして、その整数列の最小値と最大値の差が１なので、21 を 4 で割った数字 5 が基本的な数字になります。

このことから、最小値と最大値の差が１の整数列は次のいずれかになります。

「21 ÷ 4 = 5 余り 1」なので、「５」が３つと「６」が１つです。

これで最終的な整数列が求まりました。

どの数字を最終的な数字にするか

整数列 A ＝（4, 7, 3, 7）について、最終的には「5」 が３つと「６」が１つの整数列にすればいいことがわかりました。

次に考えるのは、どの数字を最終的な整数列のどの数字に操作していくかです。

最小の操作回数を求めるので、もとの整数列の「小さい数字は小さい数字」に、「大きい数字は大きい数字」に合わせる方がいいです。

つまり、元の整数列 A と最終的な整数列をそれぞれソートして（3, 4, 7, 7）を（5, 5, 5, 6）になるようにする操作回数を求めればいいことがわかります。

操作回数を求めることは、各数字の差の絶対値の合計を求めることに等しいです。

ただし、答えの操作回数は、「１減らして、１増やす」のが１回の操作なので、各数字の差の合計を２で割ったものが答えになることに注意が必要です。

• 3 を２増やせば 5 になる　⇒　|5-3| = 2
• 4 を１増やせば 5 になる　⇒　|5-4| = 1
• 7 を２減らせば 5 になる　⇒　|5-7| = 2
• 7 を１減らせば 6 になる　⇒　|6-7| = 1
• 答えは (2+1+2+1) ÷ 2 = 3

実装する

以上の考察から実装したのが次のコードです。

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


# 整数列の合計
sum_A = sum(A)

# 最終的な整数列の基本となる数字と余り
base_number = sum_A // N
amari = sum_A % N

# 最終的な整数列を作成
B = [base_number] * N

# 整数列の合計に足りない分を足す
for i in range(amari):
    B[i] += 1


# A と B をソート
A.sort()
B.sort()


answer = 0

# 各数字の差の絶対値を合計する
for i in range(N):
    answer += abs(A[i] - B[i]) 
    
# 答えの出力
print(answer//2)