python - 各要素が最大値を持つ固定長整数パーティションの一意の順列

Question

この質問は、私が数ヶ月前に持っていた質問に似ています: Generating a numpy array with all numbers of numbers that sum to less than a given number . その質問では、各要素に特定の最大値があることを考えると、最大でも定数になるすべての数値を生成したいと考えていました。

今回は、合計すると正確にその定数になるすべての順列を計算したいと思います。これは、各要素が特定の最大値を持つ整数パーティションの一意の順列を計算することと見なすことができます。最終結果は、numpy 配列に格納する必要があります。

ジェネレーターを使用すると、ワンライナーで目的が達成されます。

import numpy as np
from itertools import product
K = 3 
maxRange = np.array([1,3,2]) 

states = np.array([i for i in product(*(range(i+1) for i in maxRange)) if sum(i)==K]) 

与える

array([[0, 1, 2],
       [0, 2, 1],
       [0, 3, 0],
       [1, 0, 2],
       [1, 1, 1],
       [1, 2, 0]])

K=20と の場合、パフォーマンスが非常に遅くなりmaxRange = [20]*6ます。順列の数は 53130 と制限されていますが、既に 20 秒かかります。私の直感では、これには 1 秒もかからないはずです。

より高速なソリューションを利用できる人はいますか? これを説明するために以前の質問の解決策を変更するのに苦労していますK.

numbaの演算子を使用するソリューションは気にしません@jit...それらが現在持っているものよりも高速である限り!

前もって感謝します。