Pythonでキャンディーを一列に並んだ人々に配布するアルゴリズム

n人の人々が一列に並んでいるとき、次のようなルールに従ってキャンディーを配りたいとします。

• まず、1番目の人に1個、2番目の人に2個……というように順番に個数を増やしながら配り、最後のn番目の人にはn個渡します。
• その後、列の先頭に戻り、1番目の人にn+1個、2番目の人にn+2個……と同じ要領で配り、最後のn番目の人には2×n個渡します。

このプロセスをキャンディーがなくなるまで繰り返します。途中でキャンディーが足りなくなった場合は、その時点で残っているすべてのキャンディーをその人に渡します(直前の人より必ずしも1個多いとは限りません)。

最終的な配布結果を表す配列を返す必要があります。例えば、キャンディーの総数が7個、人数が3人の場合、出力は[2, 2, 3]となります。処理の流れを順に見てみましょう。

• 最初に1番目の人が1個もらい、配列は[1, 0, 0]になります。
• 次に2番目の人が2個もらい、配列は[1, 2, 0]になります。
• 続いて3番目の人が3個もらい、配列は[1, 2, 3]になります。
• 最後に再び先頭に戻り、1番目の人が残りの1個を受け取るので、配列は[2, 2, 3]になります。

解決のための手順

• n要素の配列resを作成し、すべて0で初期化します。
• indexを0に設定します。
• candiesが0より大きい間、以下を繰り返します。
  • res[index mod n]にmin(candies, index + 1)を加算します。
  • candiesから(index + 1)を引きます。
  • indexを1増やします。
• resを返します。

実装例

理解を深めるために、以下のPythonコードで実際の実装を確認してみましょう。

class Solution(object):
    def distributeCandies(self, candies, num_people):
        res = [0 for i in range(num_people)]
        index = 0
        while candies>0:
            res[index%num_people] += min(candies,index+1)
            candies-=(index+1)
            index+=1
        return res

ob1 = Solution()
print(ob1.distributeCandies(8, 3))

入力

8
3

出力

[3, 2, 3]

この例では、キャンディー8個を3人に配布しています。1巡目で各人が1個・2個・3個ずつ受け取り[1, 2, 3]となった後、2巡目で1番目の人が4個もらおうとしますが、残りは2個しかないため2個だけを受け取ります。そのため、最終的な結果は[3, 2, 3]となるのです。