Home > Back-end >  product of different length list using itertools in Python
product of different length list using itertools in Python

Time:01-21

I have the following lists:

sectors = ["A", "B"]

rows = [['1', '2', '3'], ['1', '2', '3', '4']]

seats = [['ab', 'abcd', 'ab'], ['ab', 'abcd', 'ab', 'abcd']]

and I want to create products like A1a, A1b, A2a, A2b, A2c ...

this code

  combinations = []
  for i in range(len(rows)):
    c = list(zip_longest(repeat(sectors[i], len(rows[i])), rows[i], seats[i]))
    combinations  = c

  for c in combinations:
    for x in product(*c):
       print("".join(x))

prints the desired results as A1a A1b A2a A2b A2c A2d A3a ...

Can this be solved in a better and more readable way, I am practicing itertools and it is a bit confusing for me.

CodePudding user response:

Not sure if that is what you want or if it is more elegant:

from itertools import chain, product

combinations = product(
    sectors,
    chain.from_iterable(chain.from_iterable(rows)),
    chain.from_iterable(chain.from_iterable(seats)),
)
joined_combinations = map(lambda t: "".join(t), combinations)
list(joined_combinations)
# returns
['A1a', 'A1b', 'A1a', 'A1b', 'A1c', 'A1d', 'A1a', ...]

Explanation: Applying two times chain.from_iterable you can "unpack" individual characters from the nested lists, then creating the product of the items of the unnested lists (which creates 3-tuples) and finally join the items of each 3-tuple together.
If you want to avoid duplicates you can put a set() around each argument in the product.

CodePudding user response:

Since you're creating triples from 3 nested lists, not Cartesian product, I think it may be more intuitive to zip the corresponding sublists and use a loop.

out = [sector   row   s for sector, rws, sts in zip(sectors, rows, seats) 
       for row, seat in zip(rws, sts) for s in seat] 

Output:

['A1a', 'A1b', 'A2a', 'A2b', 'A2c', 'A2d', 'A3a', 'A3b', 'B1a', 'B1b', 'B2a', 'B2b', 'B2c', 'B2d', 'B3a', 'B3b', 'B4a', 'B4b', 'B4c', 'B4d']

      
  •  Tags:  
  • Related