Pythonで新鮮なドーナツを受け取れるグループの最大数を求めるプログラム

問題の概要

整数 batchSize(バッチサイズ)と配列 groups が与えられます。groups[i] は、i 番目のグループに groups[i] 人の顧客がいることを意味します。あるドーナツ店では、指定された batchSize 個ずつドーナツを焼いており、「前のバッチのドーナツをすべて提供し終えるまで、次のバッチのドーナツを提供してはならない」というルールがあります。各顧客は必ずドーナツを1個受け取り、あるグループが来店した場合は、そのグループ全員の対応が終わるまで次のグループには対応できません。

グループの全員が新鮮なドーナツを受け取れたとき、そのグループは「幸せ(happy)」とみなされます。言い換えれば、グループの先頭の顧客が、直前のグループで余ったドーナツを受け取ってしまった場合、そのグループは幸せではありません。

グループの順番は自由に入れ替えられるものとし、入れ替え後に幸せになれるグループ数の最大値を求めます。

入力例

batchSize = 4、groups = [2,1,8,4,3] の場合、出力は 4 になります。グループを [8,4,2,3,1] の順に並べ替えると、1番目・2番目・3番目・4番目のグループが幸せになります。具体的には、1番目のグループ用に2バッチ、2番目のグループ用に1バッチを焼き、その後1バッチ焼いて3番目のグループへ、さらに1バッチ焼いて4番目のグループへ提供します。

アプローチ

この問題を解くには、以下の手順に従います。

• l: groups の各要素 g を batchSize で割った余り(g % batchSize)のリストを作成する
• count: l の各要素の出現回数を記録したマップを作成する
• g: 余り 0 〜 batchSize-1 それぞれの個数(count[i])のリストを作成する
• 関数 dp(sm, t) を定義する(sm は現在の累積余り、t は残っている各余りクラスのグループ数)
• t の最大値が 0 の場合(すべてのグループを処理済み)、0 を返す
• ans = 0、arr = t として初期化する
• k を 0 〜 batchSize-1 までループする:
  • arr[k] が 0 なら次の反復へスキップする
  • arr[k] を 1 減らす
  • ans = max(ans, dp((sm + k) % batchSize, arr)) で再帰的に探索する
  • arr[k] を元に戻す(バックトラック)
• sm が 0 の場合は ans + 1、そうでなければ ans を返す
• メイン処理では dp(0, g) の結果を返す

考え方のポイント

重要なのは、グループの人数そのものではなく「batchSize で割った余り」だけが結果に影響するという点です。バッチの境界でドーナツが余るかどうかは、それまでの累積人数を batchSize で割った余りだけで決まるためです。dp 関数は、次にどの余りクラスのグループを並べるかを全通り試しながら探索を行い、新しいグループの開始時点で累積和が 0(= バッチの境目と一致)であれば、そのグループは新鮮なドーナツを受け取れるとして答えに 1 を加算します。

なお、この実装はすべての並べ方を試すため、グループ数が多いほど計算量が増大します。競技プログラミング(LeetCode 1815 など)では batchSize ≤ 9 という制約のもとで、functools.lru_cache などによるメモ化を併用すると、同じ状態の再計算を避けて効率よく解くことができます。

実装例(Python)

以下の実装を見ると、理解がより深まります。

from collections import Counter

def solve(batchSize, groups):
    l = [g % batchSize for g in groups]
    count = Counter(l)
    g = [count[i] for i in range(batchSize)]

    def dp(sm, t):
        if max(t) == 0:
            return 0

        ans, arr = 0, list(t)
        for k in range(batchSize):
            if arr[k] == 0:
                continue
            arr[k] -= 1
            ans = max(ans, dp((sm + k) % batchSize, arr))
            arr[k] += 1
        return ans + (sm == 0)

    return dp(0, g)

batchSize = 4
groups = [2, 1, 8, 4, 3]
print(solve(batchSize, groups))

入力

4, [2,1,8,4,3]

出力

4