Pythonでn個のキャンディーをk個のバッグに配布する組み合わせの数を求めるプログラム

n個のキャンディーと、それらを入れるためのk個のバッグがあるとします。このとき、各バッグに必ず1個以上のキャンディーが入るように配布する方法が何通りあるかを求める問題です。ここでのポイントは、すべてのキャンディーが互いに異なる（ユニークな）ものであるという点です。そのため、どのキャンディーをどのバッグに入れるかという組み合わせをすべて数え上げる必要があります。

例えば、入力が n = 3、k = 2 の場合、出力は 3 になります。

具体的には、キャンディーは次の3通りの方法で配分できます。

(1, 2), (3)
(1), (2, 3)
(2), (1, 3)

解法のアプローチ：動的計画法（DP）

この問題は、動的計画法を使うことで効率的に解くことができます。考え方の手順は以下のとおりです。

• dp := サイズ n × n の行列を用意し、すべての要素を 1 で初期化する

• c を 2 から n まで繰り返す：

  • b を 1 から min(c, k) まで繰り返す：

    • dp[c, b] := dp[c-1, b-1] + dp[c-1, b] × (b+1)

• dp[n-1, k-1] を結果として返す

この漸化式の意味を整理してみましょう。c番目のキャンディーを処理するとき、次の2つのケースが考えられます。

• 新しいバッグに入れる場合： 残りのキャンディーで b-1 個のバッグを埋める方法は dp[c-1, b-1] 通り

• すでにあるバッグに入れる場合： 既存の b+1 個のバッグのいずれかを選べるので、dp[c-1, b] × (b+1) 通り

この2つを足し合わせることで、全体の配布方法の数が求まります。

実装例

以下のPythonコードで、実際の動作を確認できます。

def solve(n, k):
   dp = [[1] * n for _ in range(n)]
   for c in range(2, n):
      for b in range(1,min(c,k)):
         dp[c][b] = dp[c-1][b-1] + dp[c-1][b] * (b+1)
   return dp[n-1][k-1]

print(solve(3, 2))

入力

3, 2

出力

3