Search code examples
pythonarraysnumpyscipypython-itertools

Combinatoric / cartesian product of Numpy arrays without iterators and/or loop(s)


The following code

import numpy as np
import itertools

a_p1 = np.arange(0, 4, 1)
a_p2 = np.arange(20, 25, 1)

params = itertools.product(a_p1, a_p2)
for (p1, p2) in params:
    print(p1, p2)

outputs

(0, 20) (0, 21) (0, 22) (0, 23) (0, 24) (1, 20) (1, 21) (1, 22) (1, 23) (1, 24) (2, 20) (2, 21) (2, 22) (2, 23) (2, 24) (3, 20) (3, 21) (3, 22) (3, 23) (3, 24)

2 nested for loops can also outputs same results

for i, p1 in enumerate(a_p1):
    for j, p2 in enumerate(a_p2):
        print(p1, p2)

I'm looking for a solution to directly output a Numpy array with such combination (a Numpy array of tuples).

Is there a way to generate such a Numpy array without iterators and/or for loop(s) ?

I'm aware that such a solution will be more memory consuming than using iterators.


Solution

  • Install Scikit-Learn http://scikit-learn.org/

    from sklearn.utils.extmath import cartesian
    print cartesian([a_p1, a_p2])
    

    It should output

    [[ 0 20]
     [ 0 21]
     [ 0 22]
     [ 0 23]
     [ 0 24]
     [ 1 20]
     [ 1 21]
     [ 1 22]
     [ 1 23]
     [ 1 24]
     [ 2 20]
     [ 2 21]
     [ 2 22]
     [ 2 23]
     [ 2 24]
     [ 3 20]
     [ 3 21]
     [ 3 22]
     [ 3 23]
     [ 3 24]]
    

    This solution was taken from a similar question: Using numpy to build an array of all combinations of two arrays