I have a problem finding all combinations of a 2D array. Let’s suggest I have an array as follwoing:
[ [1,2], [3,4], [5,6] ]
Now I need to get all possible combninations such as
[ [1,3,5], [1,3,6], [1,4,5], [1,4,6], [2,3,5], [2,3,6], [2,4,5], [2,4,6] ]
I’ve managed to get it with numpy and meshgrids with single arrays:
import numpy as np array_1 = np.array([1,2]) array_2 = np.array([3,4]) array_3 = np.array([5,6]) combination = np.array(np.meshgrid(array_1, array_2, array_3)).T.reshape(-1,3)
But with a 2D array I cant’t get it to work.
import numpy as np multi = [np.array([1,2]), np.array([3,4]), np.array([5,6])] combination = np.array(np.meshgrid(multi)).T.reshape(-1,len(multi))
This doesn’t give me the expected result but only:
[ [1,2,3], [4,5,6] ]
Is there any way to get this working correctly?
Advertisement
Answer
You can use the itertools
package in the standard library.
itertools.product
generates all combinations that you wish.
import itertools arrays = [ [1,2], [3,4], [5,6] ] list(itertools.product(*arrays)) #[(1, 3, 5), # (1, 3, 6), # (1, 4, 5), # (1, 4, 6), # (2, 3, 5), # (2, 3, 6), # (2, 4, 5), # (2, 4, 6)]
If you prefer lists to tuples, you can do the following:
list(map(list, itertools.product(*arrays))) #[[1, 3, 5], # [1, 3, 6], # [1, 4, 5], # [1, 4, 6], # [2, 3, 5], # [2, 3, 6], # [2, 4, 5], # [2, 4, 6]]