ABC208 C問題 Fair Candy Distribution を直感で解く - 競技プログラミング悪戦苦闘日記

ABC208 C問題をゆるく解説します。

問題

高橋王国には N 人の国民がいます。全ての国民には国民番号が割り振られており、 i 人目の国民の国民番号は $a_i$ です。ここで、 $a_i$ は互いに異なります。
高橋君は K 個のお菓子を持っています。高橋君は次のルールに従って、持っているお菓子が無くなるまで国民にお菓子を配ることにしました。
持っているお菓子が N 個以上ある場合、全員に 1 個ずつお菓子を配る。
そうでない場合、その時点で高橋くんが持っているお菓子の個数を K′ として、国民番号が小さい方から K′ 人に 1 個ずつ配る。より厳密には、 $a_i$ の値が小さい方から K′ 人を選び、選んだ人に 1 個ずつお菓子を配る。
全てのお菓子を配り終えたとき、i 人目の国民は何個のお菓子を持っていますか？

難しいポイント

Nの最大値が $2 \times 10^5$ 、Kの最大値が $10^{18}$ なので、愚直に解けないことがわかります。
一番難しいのは問題を理解することなのかなとも思います。

解き方の方針

サンプル１の入力例を愚直に手で書き出して解法の糸口が見えました。
$N = 2, K = 7, a_1 = 1, a_2 = 8$ の場合を考えると持っているお菓子の数が $N$ を下回るまで、全員にお菓子を1個ずつ配ってます。
お菓子の残りが1個になっているということは、6個のお菓子を2人に等しく配っていることになります。
つまり、 $7 \div 2$ の商の3個は全員に行き渡ります。
一般化すると、 $\lfloor \frac{K}{N} \rfloor$ 、KをNで割った切り捨てた数が少なくとも全員に行き渡ります。
あとはKをNで割った余りをどう分配させるかです。
このサンプルでは余りは1です。
問題文から $a_i$ の値が小さい方から配っていくということなので、小さい方から1つまでは配られることがわかります。
つまりKをNで割った余りの数をrとすると、小さい方からr番目までは全員に行き渡る個数（ $\lfloor \frac{K}{N} \rfloor$ ）より1つ大きい数になります。

実装

公式解説の実装の方が美しいです。
私は最初の順番と $a_i$ の値を組み合わせた配列（1つ目を順番、2つ目を $a_i$ ）を配列に格納して、2次元配列にしました。
まずは、2次元配列の2つ目の要素である $a_i$ の値で小さい順にソートして、
$\frac{K}{N}$ の余りをrとした時に、最初のr番目までは、最低限全員に行き渡る個数 $\lfloor \frac{K}{N} \rfloor$ に1を足した数を格納し、それより後は最低限全員に行き渡る個数を格納しました。
最後に2次元配列の1つ目の要素でソートして、先頭から順番に2つ目の要素の値を出力します。

N, K = map(int, input().split())
a_arr = list(map(int, input().split()))
amap_arr = []
for i, a in enumerate(a_arr):
    amap_arr.append([i, a])
# 2つ目の要素a_iでソート
amap_arr = sorted(amap_arr, key=lambda x: x[1])
# 最低限全員に配られる個数
at_least = K // N
amari = K % N
for i in range(N):
    # 最初のamari個はat_leastに1を足した個数がもらえる
    if i <= amari - 1:
        amap_arr[i][1] = at_least + 1
        continue
    amap_arr[i][1] = at_least

# 最初の順番になるようにソート
amap_arr = sorted(amap_arr, key=lambda x: x[0])
# 答えを出力
for i in range(N):
    print(amap_arr[i][1])