I need to make a combination of a list of lists, selecting n elements o each list for example
a=[[1,2,3,4,5],[6,7,8,9,10]]
n1=2
n2=3
so my result can be something like this:
r=[[1,2,6,7,8],[1,2,6,7,9],...,[4,5,7,8,9],[4,5,8,9,10]]
Is there any clean way to do it? Or should I have to break my lists into smaller sizes and use for loops to call the itertools?
Simply generate the combinations of the two lists separately, then take the Cartesian product of the two generators:
from itertools import product, combinations
r_gen = product(combinations(a[0], n1), combinations(a[1], n2))
r = (a + b for a, b in r_gen)
The first 10 elements yielded by r
are
[(1, 2, 6, 7, 8),
(1, 2, 6, 7, 9),
(1, 2, 6, 7, 10),
(1, 2, 6, 8, 9),
(1, 2, 6, 8, 10),
(1, 2, 6, 9, 10),
(1, 2, 7, 8, 9),
(1, 2, 7, 8, 10),
(1, 2, 7, 9, 10),
(1, 2, 8, 9, 10)]
If I understood correctly, there are basically two steps in this problem:
itertools.combinations
(or .permutations
based on what you want:a1 = itertools.combinations(a[0], n1)
a2 = itertools.combinations(a[1], n2)
r = itertools.product(a1, a2)
To make the result look exactly what you are looking for, you can use a list comprehension to concatenate the tuples:
r = [list(s1 + s2) for s1, s2 in r
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.