Hay varios ejemplos elegantes del uso de numpy en Python para generar matrices de todas las combinaciones. Por ejemplo, la respuesta aquí:Usando numpy para construir una matriz de todas las combinaciones de dos matrices .

Ahora suponga que hay una restricción adicional, a saber, la suma de todos los números no puede sumar más de una constante dadaK. Usando un generador yitertools.product, por ejemplo conK=3 donde queremos las combinaciones de tres variables con rangos 0-1,0-3 y 0-2 podemos hacerlo de la siguiente manera:

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])

que vuelve

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

En principio, el enfoque dehttps://stackoverflow.com/a/25655090/1479342 se puede usar para generar todas las combinaciones posibles sin la restricción y luego seleccionar el subconjunto de combinaciones que suman menos deK. Sin embargo, ese enfoque genera muchas más combinaciones de las necesarias, especialmente siK es relativamente pequeño en comparación consum(maxRange).

Debe haber una manera de hacerlo más rápido y con un menor uso de memoria. ¿Cómo se puede lograr esto usando un enfoque vectorizado (por ejemplo, usandonp.indices)?