重複組み合わせ

問題（重複組み合わせ）

$ n $ 種類の品物があり、 $ i $ 番目の品物は $ a_i $ 個あります。

異なる種類の品物同士は区別できますが、同じ種類の品物同士は区別できません。

これらの品物の中から $ m $ 個選ぶ組み合わせの総数を求めて下さい。

[制約]
🔹$ 1 \leqq n \leqq 1000 $
🔹$ 1 \leqq m \leqq 1000 $
🔹$ 1 \leqq a_i \leqq 1000 $

解き方・ソースコード

重複なく数え上げるために、同じ種類の品物を一度に処理します。

漸化式を次のように定義します。

$$ dp[i+1][j] := i 番目までの品物からj個選ぶ組み合わせの総数 $$

$ i $ 番目までの品物から $ j $ 個選ぶためには、$ i - 1 $ 番目までの品物から $ j - k $ 個選んで、$ i $ 番目の品物を $ k $ 加えればよいので

$$ dp[i+1][j] = \sum_{k=0}^{min(j,a[i])}dp[i][j-k] $$

という漸化式が成り立ちます。

この漸化式は計算量が $ O(nm^2) $ となりますが、

$$ \sum_{k=0}^{min(j,a[i])}dp[i][j-k] = \sum_{k=0}^{min(j-1,a[i])}dp[i][j-1-k]+dp[i][j]-dp[i][j-1-a_i] $$

であるので、

$$ dp[i+1][j]=dp[i+1][j-1] + dp[i][j]-dp[i][j-1-a_i] $$

と変形することができ、計算量は$ O(nm) $ となります。

[Google Colaboratory]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
#--------- 入力例 ----------
n = 3           # 品物の数
a = [1, 2, 3]   # 各品物の個数
m = 3           # 分割数
#---------------------------
memo = {}

# １つも選ばない方法は常に一通り
for i in range(n + 1):
    memo[i, 0] = 1

for i in range(n):              # 品物の数に対応するループ
    for j in range(1, m + 1):   # 分割数に対応するループ
        if j - i -a[i] >= 0:
            memo[i + 1, j] = memo.get((i + 1, j - 1), 0) + memo.get((i, j), 0) - memo.get((i, j - 1 - a[i]), 0)
        else:
            memo[i + 1, j] = memo.get((i + 1, j - 1), 0) + memo.get((i, j), 0)

print('解：', memo[m, n])

[実行結果]

解： 6

[1 2 0]、[1 1 1]、[1 0 2]、[0 2 1]、[0 1 2]、[0 0 3]の６パターンあるので、解は6となります。