Home > Mobile >  Combinatorics in python numpy
Combinatorics in python numpy

Time:04-09

I'd like to write a single function that outputs all the possible combinations of 2 matrixes:

def combine(*args):
    return np.array(np.meshgrid(args)).T.reshape(-1, len(args) 1)

However when passed:

print(combine(np.array([1,2,3]), np.array([4,5,6])))

It outputs:

[[1 2 3]
 [4 5 6]]

How can I make it work? I would like to keep it automatic, not to simply pass (args[0], args[1])

CodePudding user response:

A straight forward use of python itertools:

In [134]: import itertools
In [135]: a,b = [1,2,3], [4,5,6]
In [137]: list(itertools.product(a,b))
Out[137]: [(1, 4), (1, 5), (1, 6), (2, 4), (2, 5), (2, 6), (3, 4), (3, 5), (3, 6)]

or as array:

In [145]: np.array(list(itertools.product(a,b)))
Out[145]: 
array([[1, 4],
       [1, 5],
       [1, 6],
       [2, 4],
       [2, 5],
       [2, 6],
       [3, 4],
       [3, 5],
       [3, 6]])
  • Related