CODEWITHSUNDEEP
pythoncombinationspython-itertools
Marcelo Ruiz
Marcelo Ruiz

Reputation: 393

Itertools combinations of multiple list selecting n elements per list

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?

Upvotes: 1

Views: 1122

Answers (2)

M. Zhang
M. Zhang

Reputation: 910

If I understood correctly, there are basically two steps in this problem:

  1. Choose n items from each group. This can be done with itertools.combinations (or .permutations based on what you want:
a1 = itertools.combinations(a[0], n1)
a2 = itertools.combinations(a[1], n2)
  1. Find the combinations of those two iterables. This is almost what the Cartesian product does:
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

Upvotes: 1

Brian61354270
Brian61354270

Reputation: 14423

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)]

Upvotes: 5

Related Questions