numpy array of combination of number solution?

user_16922127 0
Viewed 71

I am trying to create a numpy array of the following type:

range = 3, size = 4    => [(0,0,0,0), (0,0,0,1) .... (0,0,0,3), (0,0,1,0)........(3,3,3,3)]

currently, I have got the solution of using itertool.combination to generate the list and then convert it into numpy array.

Is there any smarter and better / faster way to implement it? especially when the range and the size number both are going up to 50?

numpy python
1 Answers

Repeated broadcast assignment can achieve:

def combi(rng, size):
    rng += 1
    ar = np.empty((rng ** size, size), dtype=int)
    a = np.arange(rng)
    shape = (rng,) * size
    for col in ar.T[::-1].reshape(-1, *shape):
        col[...] = a
        a = a[..., None]
    return ar

Test:

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

But for your requirements, the shape of the generated array will reach (51 ** 50, 50). I think no machine can bear such an array.

user_179809310
Related