abudis
abudis

Reputation: 2881

Create a list of lists of possible combinations

I have a list, which contains the following items:

initl = [1, 2, 3, 4]

and a stepsize, which is stepsize = 0.05. Each value in the list is allowed to change by the stepsize either up or down. What I would like to create is a list of lists, which contains all combinations of up, down or initial values, e.g. like so:

result_list = [[1, 2, 3, 4], [1.05, 2, 3, 4], [0.95, 2, 3, 4], [1, 2.05, 3, 4], ...] 

The order of combinations within the list doesn't matter.

I came up with this:

import itertools


initl = [1, 2, 3, 4]
stepsize = 0.05

lu = [i + stepsize for i in initl]
ld = [i - stepsize for i in initl]

l_map = list(itertools.product(range(3), repeat=4))

result_list = []

for i in l_map:
    l_new = []
    for pos, j in enumerate(i):
        if j == 0:
            l_new.append(ld[pos])
        elif j == 1:
            l_new.append(initl[pos])
        else:
            l_new.append(lu[pos])
   result_list.append(l_new)

which produces the desired output. List length: 3^4 = 81. But I wonder if there's a better approach to this. The nested for loop in particular seems clunky to me here. Any help appreciated.

Upvotes: 1

Views: 69

Answers (4)

CristiFati
CristiFati

Reputation: 41116

You're almost there. All you need to do is transposing the matrix consisting of the 3 (4 lengthed) lists, so that itertools.product can pick for each element from [elem[i] - stepsize, elem[i], elem[i] + stepsize]. That is done by the zip function.

>>> import itertools as it
>>>
>>> initl = [1, 2, 3, 4]
>>> stepsize = 0.05
>>>
>>> lo = [elem - stepsize for elem in initl]
>>> hi = [elem + stepsize for elem in initl]
>>> lo, hi
([0.95, 1.95, 2.95, 3.95], [1.05, 2.05, 3.05, 4.05])
>>>
>>> list(it.product(*zip(lo, initl, hi)))
[(0.95, 1.95, 2.95, 3.95), (0.95, 1.95, 2.95, 4), (0.95, 1.95, 2.95, 4.05), (0.95, 1.95, 3, 3.95), (0.95, 1.95, 3, 4), (0.95, 1.95, 3, 4.05), (0.95, 1.95, 3.05, 3.95), (0.95, 1.95, 3.05, 4), (0.95, 1.95, 3.05, 4.05), (0.95, 2, 2.95, 3.95), (0.95, 2, 2.95, 4), (0.95, 2, 2.95, 4.05), (0.95, 2, 3, 3.95), (0.95, 2, 3, 4), (0.95, 2, 3, 4.05), (0.95, 2, 3.05, 3.95), (0.95, 2, 3.05, 4), (0.95, 2, 3.05, 4.05), (0.95, 2.05, 2.95, 3.95), (0.95, 2.05, 2.95, 4), (0.95, 2.05, 2.95, 4.05), (0.95, 2.05, 3, 3.95), (0.95, 2.05, 3, 4), (0.95, 2.05, 3, 4.05), (0.95, 2.05, 3.05, 3.95), (0.95, 2.05, 3.05, 4), (0.95, 2.05, 3.05, 4.05), (1, 1.95, 2.95, 3.95), (1, 1.95, 2.95, 4), (1, 1.95, 2.95, 4.05), (1, 1.95, 3, 3.95), (1, 1.95, 3, 4), (1, 1.95, 3, 4.05), (1, 1.95, 3.05, 3.95), (1, 1.95, 3.05, 4), (1, 1.95, 3.05, 4.05), (1, 2, 2.95, 3.95), (1, 2, 2.95, 4), (1, 2, 2.95, 4.05), (1, 2, 3, 3.95), (1, 2, 3, 4), (1, 2, 3, 4.05), (1, 2, 3.05, 3.95), (1, 2, 3.05, 4), (1, 2, 3.05, 4.05), (1, 2.05, 2.95, 3.95), (1, 2.05, 2.95, 4), (1, 2.05, 2.95, 4.05), (1, 2.05, 3, 3.95), (1, 2.05, 3, 4), (1, 2.05, 3, 4.05), (1, 2.05, 3.05, 3.95), (1, 2.05, 3.05, 4), (1, 2.05, 3.05, 4.05), (1.05, 1.95, 2.95, 3.95), (1.05, 1.95, 2.95, 4), (1.05, 1.95, 2.95, 4.05), (1.05, 1.95, 3, 3.95), (1.05, 1.95, 3, 4), (1.05, 1.95, 3, 4.05), (1.05, 1.95, 3.05, 3.95), (1.05, 1.95, 3.05, 4), (1.05, 1.95, 3.05, 4.05), (1.05, 2, 2.95, 3.95), (1.05, 2, 2.95, 4), (1.05, 2, 2.95, 4.05), (1.05, 2, 3, 3.95), (1.05, 2, 3, 4), (1.05, 2, 3, 4.05), (1.05, 2, 3.05, 3.95), (1.05, 2, 3.05, 4), (1.05, 2, 3.05, 4.05), (1.05, 2.05, 2.95, 3.95), (1.05, 2.05, 2.95, 4), (1.05, 2.05, 2.95, 4.05), (1.05, 2.05, 3, 3.95), (1.05, 2.05, 3, 4), (1.05, 2.05, 3, 4.05), (1.05, 2.05, 3.05, 3.95), (1.05, 2.05, 3.05, 4), (1.05, 2.05, 3.05, 4.05)]

Upvotes: 1

Georgina Skibinski
Georgina Skibinski

Reputation: 13377

Try:

def getComb(arr, step, res=[]):
    if(len(arr)==1):
        for el in [-step, 0, step]:
            yield res+[arr[0]+el]
    else:
        for el in [-step, 0, step]:
            yield from getComb(arr[1:], step, res+[arr[0]+el])

For your input data outputs:

>>> for el in getComb(initl, stepsize): print(el)

[0.95, 1.95, 2.95, 3.95]
[0.95, 1.95, 2.95, 4]
[0.95, 1.95, 2.95, 4.05]
[0.95, 1.95, 3, 3.95]
[0.95, 1.95, 3, 4]
...
[1.05, 2.05, 3, 4]
[1.05, 2.05, 3, 4.05]
[1.05, 2.05, 3.05, 3.95]
[1.05, 2.05, 3.05, 4]
[1.05, 2.05, 3.05, 4.05]

Upvotes: 1

tobias_k
tobias_k

Reputation: 82889

You could try a list comprehension like the following, ziping the products with initl:

>>> from itertools import product
>>> initl = [1, 2, 3, 4]
>>> step = [-0.05, 0, 0.05]
>>> [[x+d for x, d in zip(initl, p)] for p in product(step, repeat=len(initl))]
[[0.95, 1.95, 2.95, 3.95],
 [0.95, 1.95, 2.95, 4],
 ...
 [1.05, 2.05, 3.05, 4.05]]
>>> len(_)
81

Upvotes: 1

Adam.Er8
Adam.Er8

Reputation: 13393

You had the right idea, essentially implementing numpy.choose yourself, for a 3-choice case.

You can simplify using itertools.product directly on a list of which contains a tuple of all options for each value, like this:

import itertools

initl = [1, 2, 3, 4]
stepsize = 0.05

prod_seed = [(i, i+stepsize, i-stepsize) for i in initl]
result_list = list(itertools.product(*prod_seed))

print(result_list)
print(len(result_list))

Output:

[(1, 2, 3, 4), (1, 2, 3, 4.05), (1, 2, 3, 3.95), ....]
81

Upvotes: 1

Related Questions