pythonarraysnumpygeneralization

Pairwise combination of arbitrary set of lists


I am looking for a generalized approach that solves the following problem: Given an arbitrary number of one- or multidimensional lists (NumPy arrays or the like) return their pairwise combinations.

The following tries to give an example for three one-dimensional lists a, b and c that result in d:

a = np.linspace(0, 1, 1).astype(int)
b = np.linspace(1, 2, 2).astype(int)
c = np.linspace(2, 4, 3).astype(int)

e = np.array([
    (a_, b_, c_)
    for a_ in a
    for b_ in b
    for c_ in c
])

When executed the variables are set as follows:

# a
[0]

# b
[1 2]

# c
[2 3 4]

# d
[[0 1 2]
 [0 1 3]
 [0 1 4]
 [0 2 2]
 [0 2 3]
 [0 2 4]]

What would be a good way for a generalized approach? I'm ideally looking for a function that takes an iterable whose elements define the one- or multidimensional lists like so:

def pairwise_combinations(iterable):
    # Insert magic here

Solution

  • itertools.product got your back: https://docs.python.org/3/library/itertools.html#itertools.product

    In the given example it could be used like this:

    d = np.array(list(itertools.product(a, b, c)))