Skip to content
Advertisement

All possible combinations of arrays in python

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]]
User contributions licensed under: CC BY-SA
5 People found this is helpful
Advertisement