There are several elegant examples of using numpy in Python to generate arrays of all combinations. For example the answer here: Using numpy to build an array of all combinations of two arrays .
Now suppose there is an extra constraint, namely, the sum of all numbers cannot add up to more than a given constant K
. Using a generator and itertools.product
, for an example with K=3
where we want the combinations of three variables with ranges 0-1,0-3, and 0-2 we can do it a follows:
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])
which returns
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]])
In principle, the approach from https://stackoverflow.com/a/25655090/1479342 can be used to generate all possible combinations without the constraint and then selecting the subset of combinations that sum up to less than K
. However, that approach generates much more combinations than necessary, especially if K
is relatively small compared to sum(maxRange)
.
There must be a way to do this faster and with lower memory usage. How can this be achieved using a vectorized approach (for example using np.indices
)?