CODEWITHSUNDEEP
pythonarrayscombinationsnested-loops
nkz
nkz

Reputation: 165

[Python]: generate and array of all possible combinations

I have a very straightforward combination problem. I have two arrays (a and b). Array a indicates all the values one of the three slots in array b can take on. Each slot in array b can have a value between 1 and 5. An example of which would be [1, 4, 5]. I would like to generate an array (c) with all possible combinations. I like to extend the basic example larger arrays.

Input: 

a = [1, 2, 3, 4, 5]
b = [1, 2, 3]

Output:

c = [[1, 1, 1], [1, 1, 2],[1, 1, 3], [1, 1, 4], [1, 1, 5],
     [1, 2, 1], [1, 2, 2],[1, 2, 3], [1, 2, 4], [1, 2, 5],
     [1, 3, 1], [1, 3, 2],[1, 3, 3], [1, 3, 4], [1, 3, 5],
     [1, 4, 1], [1, 4, 2],[1, 4, 3], [1, 4, 4], [1, 4, 5],
     [1, 5, 1], [1, 5, 2],[1, 5, 3], [1, 5, 4], [1, 5, 5],
     [2, 1, 1], [2, 1, 2],[2, 1, 3], [2, 1, 4], [2, 1, 5],
     [2, 2, 1], [2, 2, 2],[2, 2, 3], [2, 2, 4], [2, 2, 5],
     [2, 3, 1], [2, 3, 2],[2, 3, 3], [2, 3, 4], [2, 3, 5],
     [2, 4, 1], [2, 4, 2],[2, 4, 3], [2, 4, 4], [2, 4, 5],
     [2, 5, 1], [2, 5, 2],[2, 5, 3], [2, 5, 4], [2, 5, 5],
     [3, 1, 1], [3, 1, 2],[3, 1, 3], [3, 1, 4], [3, 1, 5],
     [3, 2, 1], [3, 2, 2],[3, 2, 3], [3, 2, 4], [3, 2, 5],
     [3, 3, 1], [3, 3, 2],[3, 3, 3], [3, 3, 4], [3, 3, 5],
     [3, 4, 1], [3, 4, 2],[3, 4, 3], [3, 4, 4], [3, 4, 5],
     [3, 5, 1], [3, 5, 2],[3, 5, 3], [3, 5, 4], [3, 5, 5],
     [4, 1, 1], [4, 1, 2],[4, 1, 3], [4, 1, 4], [4, 1, 5],
     [4, 2, 1], [4, 2, 2],[4, 2, 3], [4, 2, 4], [4, 2, 5],
     [4, 3, 1], [4, 3, 2],[4, 3, 3], [4, 3, 4], [4, 3, 5],
     [4, 4, 1], [4, 4, 2],[4, 4, 3], [4, 4, 4], [4, 4, 5],
     [5, 5, 1], [5, 5, 2],[5, 5, 3], [5, 5, 4], [5, 5, 5],
     [5, 1, 1], [5, 1, 2],[5, 1, 3], [5, 1, 4], [5, 1, 5],
     [5, 2, 1], [5, 2, 2],[5, 2, 3], [5, 2, 4], [5, 2, 5],
     [5, 3, 1], [5, 3, 2],[5, 3, 3], [5, 3, 4], [5, 3, 5],
     [5, 4, 1], [5, 4, 2],[5, 4, 3], [5, 4, 4], [5, 4, 5],
     [5, 5, 1], [5, 5, 2],[5, 5, 3], [5, 5, 4], [5, 5, 5]]

Solution for the problem above:

d = []
for i in range(len(a)):
    for j in range(len(a)):
        for k in range(len(a)):
            e = []
            e.append(i+1)
            e.append(j+1)
            e.append(k+1)
            d.append(e)

I am looking for a more generic way. One which can accommodate larger arrays (see below) without the need to use a nested for loop structure. I searched for a comparable example, but was unable to find one on stackoverflow.

Input:

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
b = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]

Upvotes: 4

Views: 2460

Answers (2)

tobias_k
tobias_k

Reputation: 82949

In the general case, you should of course use the itertools module, in this particular case itertools.product, as explained in the other answer.

If you want to implement the function yourself, you can use recursion to make it applicable to any array sizes. Also, you should probably make it a generator function (using yield instead of return), as the result could be rather long. You can try something like this:

def combinations(lst, num):
    if num > 0:
        for x in lst:
            for comb in combinations(lst, num - 1):
                yield [x] + comb
    else:
        yield []

Upvotes: 2

Delgan
Delgan

Reputation: 19697

You are looking for itertools.product().

a = [1, 2, 3, 4, 5]
b = 3  # Actually, you just need the length of the array, values do not matter

c = itertools.product(a, repeat=b)

Note that this returns an iterator, you may need to cast it using list() but be aware this can take forever and highly consume memory if the sizes grow.

Upvotes: 4

Related Questions