AtCoder が提供しているABC（AtCoder Beginner Contest）313 のC問題をC++とPythonで解いてみました。ABC313は、2023年8月5日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
C問題 Approximate Equalization 2（Difficulty : 681）
問題はリンク先をご覧ください。
ABC313 C問題 Approximate Equalization 2
操作によって数列の合計が変わらないことがポイントとなります。AtCoder Problems による Difficulty は 681 でした。
解答案
問題で与えられる操作は、要素の数を1減らして、別の要素の数を1増やします。このため、数列 A の合計は操作によって変わりません。
A の最小値と最大値の差が1以下になるときに、どの数に収束しているか考えます。
１）数列の合計が要素の数 N で割り切れる場合
この場合は、すべての要素が平均（合計をNで割った値）に収束します。平均より小さい差の合計と大きい差の合計が一致します。このため、ある平均よりひとつ小さい要素がある場合、組として平均よりひとつ大きな要素が必ず存在します。このため、すべての要素が平均に収束していないと差が1以下になりません。すべての要素が平均に収束するための回数は、平均より小さい差の合計（＝平均より大きい差の合計）となります。
例として、数列 1、1、1、9 を考えます。平均はちょうど3になります。この数列に対して、以下のように操作を加えます。
(1, 1, 1, 9) → (3, 1, 1, 7) → (3, 3, 1, 5) → (3, 3, 3, 3)
平均3に収束するために 2×3＝6 回の操作が必要となります。
２）数列の合計が要素の数 N で割り切れない場合
この場合は、平均が整数になりません（小数部を持ちます）。平均より小さい一番近い整数または平均より大きい一番近い整数と、数列の要素の差を考えます。１）の場合と異なり、平均を挟む2つの整数より小さい差の合計と大きい差の合計が一致しないことがあり得ます。この場合、差の合計の大きい方が必要な操作回数となります。
具体的な例で考えます。数列 1、1、1、11 を考えます。平均は3.5になります。このため3か4に収束します。
(1, 1, 1, 11) → (3, 1, 1, 9) → (3, 3, 1, 7) → (3, 3, 3, 5) → (3, 3, 4, 4)
3より小さい要素と差の合計が6で、4より大きい要素の差の合計が7となります。この場合は、小さい要素をすべて3に収束させても（操作6回）、条件を満たさない (3, 3, 3, 5) となります。もう1回操作を行い、(3, 3, 4, 4) になります（操作7回）。
C++ プログラム例（ABC313C）
上記の考察をプログラムに変換しました。2）の場合は、平均を挟む2つの整数より小さい差の合計（result1）と大きい差の合計（result2）を求めて、大きい値を出力しています（28ー35行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n;
cin >> n;
vector<int> a(n);
ll sum = 0;
for (int i = 0; i < n; ++i) {
cin >> a[i];
sum += a[i];
}
ll ave = sum / n;
ll result = 0;
if (sum % n == 0) {
for (int i = 0; i < n; ++i) {
if (a[i] < ave) {
result += ave - a[i];
}
}
} else {
ll result1 = 0;
ll result2 = 0;
for (int i = 0; i < n; ++i) {
if (a[i] < ave) {
result1 += ave - a[i];
} else if (a[i] > ave + 1) {
result2 += a[i] - (ave + 1);
}
}
result = max(result1, result2);
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC313C）
C++ プログラムを移植しました。変数 sum は、組込み関数と命名が重なっているため変数を s に変更しました。
"""AtCoder Beginner Contest 313 C"""
n = int(input())
a = list(map(int, input().split()))
s = sum(a)
ave = s // n
result = 0
if s % n == 0:
for i in range(n):
if a[i] < ave:
result += ave - a[i]
else:
result1 = 0
result2 = 0
for i in range(n):
if a[i] < ave:
result1 += ave - a[i]
elif a[i] > ave + 1:
result2 += a[i] - (ave + 1)
result = max(result1, result2)
print(result)こちらも「AC」と判定されました。
最後に
この問題を解くために、特定のアルゴリズムの知識が求められるわけではありません。似たような問題を解いて、自分の中で消化していくしかない気がします。
引き続き ABC の問題を紹介していきます。
