Creating the Cartesian Product of a set of Vectors in Python?

Question

Given the standard basis vectors (e_1,e_2,e_3) in 3 dimensions and letting the elements of (e_1,e_2,e_3) be restricted to, say (0,1,2,3,4) is there a simple pythonic way to create the cartesian product of all the vectors in this vector space?

For example, given [1,0,0],[0,1,0] and [0,0,1], I would like to get a list of all of the linear combinations (where the a_i's are restricted to the naturals between 0 and 4) of these vectors between [0,0,0] and [4,4,4].

I could program this up myself but before going to that trouble I thought I would ask if there is a simple pythonic way of doing it, maybe in numpy or something similar.

Answer 1

For the specific case of a space of natural numbers, you want np.indices:

>>> np.indices((4, 4)).reshape(2,-1).T
array([[0, 0],
       [0, 1],
       [0, 2],
       [0, 3],
       [1, 0],
       [1, 1],
       [1, 2],
       [1, 3],
       [2, 0],
       [2, 1],
       [2, 2],
       [2, 3],
       [3, 0],
       [3, 1],
       [3, 2],
       [3, 3]])

(numpy actually outputs these in a grid, but you wanted a 1-D list of points, hence the .reshape)

Otherwise, what you're describing is not a powerset but a cartesian product

itertools.product(range(4), repeat=3)