I have a problem with 'flattening' out some generators in python. Here is my code:
import itertools as it
test = [[1,2,3],[4,5],[6,7,8]]
def comb(possible):
if len(possible) != 1:
for a in possible[0]:
yield from it.product((a,), comb(possible[1:]))
else:
yield from possible[0]
list(comb(test))
which gives me:
[(1, (4, 6)),
(1, (4, 7)),
(1, (4, 8)),
(1, (5, 6)),
(1, (5, 7)),
(1, (5, 8)),
(2, (4, 6)),
(2, (4, 7)),
(2, (4, 8)),
(2, (5, 6)),
(2, (5, 7)),
(2, (5, 8)),
(3, (4, 6)),
(3, (4, 7)),
(3, (4, 8)),
(3, (5, 6)),
(3, (5, 7)),
(3, (5, 8))]
However, I want something like:
[(1, 4, 6),
(1, 4, 7),
(1, 4, 8),
(1, 5, 6),
(1, 5, 7),
(1, 5, 8),
(2, 4, 6),
(2, 4, 7),
(2, 4, 8),
(2, 5, 6),
(2, 5, 7),
(2, 5, 8),
(3, 4, 6),
(3, 4, 7),
(3, 4, 8),
(3, 5, 6),
(3, 5, 7),
(3, 5, 8)]
In general the function should give me generators for all possible paths to go through a list, i.e. from test[0] -> test[1] -> ... -> test[n]
where n
is len(test)
. Here, it picks up at each step one element.
Similar to what the following function returns, just with generators:
def prod(possible):
if len(possible) != 1:
b = []
for i in range(len(possible[0])):
for x in prod(possible[1:]):
if len(possible) == 2:
b += [[possible[0][i]]+[x]]
else:
b += [[possible[0][i]]+x]
return b
else:
return possible[0]
prod(test)
I played around with it.chain
and it.chain.from_iterable
but can't seem to make it work. The problem is that my 'test' list are variable in size and length and thus I have to do the whole thing recursively.
Edit:
itertools.product(*test)
works as pointed out by John Coleman