1

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?

Marcelo Ruiz
  • 373
  • 3
  • 14

2 Answers2

5

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)]
Brian61354270
  • 8,690
  • 4
  • 21
  • 43
  • I think you miss `list` in `r = list(a + b for a, b in r_gen)` `r` is still a `generator`. – Daniel Hao Nov 15 '21 at 02:10
  • Thanks Brian for your answer, if I have also a third list, a[2], n3, could I combine the same approach? e.g. using the result list r, with n_r=length of the lists: r_gen = product(combinations(r, n_r]), combinations(a[2], n3)) ? or is there a direct way to do it? – Marcelo Ruiz Nov 15 '21 at 08:33
  • @Pulse9 [`itertools.product`](https://docs.python.org/3/library/itertools.html#itertools.product) accepts an arbitrary number of iterables, so you can simplify pass it all three combinations iterators. You can also generalize this for an arbitrary number of list by using something like `product(*[combinations(lst, num) for lst, num in zip(a, n)]` where `a` is your list of lists and `n` is a list of choices-per-list. – Brian61354270 Nov 15 '21 at 15:52
  • That is perfect! Thank you very much Brian – Marcelo Ruiz Nov 15 '21 at 16:04
  • Just final question, what does the * mean in the argument? – Marcelo Ruiz Nov 15 '21 at 16:05
  • 1
    @Pulse9 The `*` signifies [argument unpacking](https://docs.python.org/3/tutorial/controlflow.html#unpacking-argument-lists). It lets you pass the elements of a iterable as separate arguments to a function, so `foo(*[a, b, c])` has the same effect as writing `foo(a, b, c)`. See also [Pass a list to a function to act as multiple arguments](https://stackoverflow.com/q/3480184/11082165). – Brian61354270 Nov 15 '21 at 16:17
1

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
M. Zhang
  • 750
  • 8
  • 18