How to split a nested list into several lists in Python?

Question

The input format of my code is as follows:

The first line contains an integer n . The next n lines contain a list of space separated integers.

What I need is to convert the elements of each line into a list and then compute the cartesian product of those lists. So I've reached the point where I convert the elements of each line into a list and store the lists in "mylist".

However, since "mylist" is a nested list, I do know how to compute the cartesian product of each element.

from itertools import product
n = int(input())

mylist = []
for i in range(n):
    elem = list(map(int, input().split()))
    mylist.append(elem)
product = list(product(???))

For example, if my input is:

2 # number of lists
1 2 3 # 1st list
4 5 6 # 2nd list

then "mylist" is gonna be:

my list = [[1, 2, 3], [4, 5, 6]]

and I need the following outupt (cartesian product of the 2 lists in "mylist"):

[(1, 4), (1, 5), (1, 6), (2, 4), (2, 5), (2, 6), (3, 4), (3, 5), (3, 6)]

OBS: I don't necessarily need a variable called "mylist"; I just need the cartesian product of the 2 input lines.

Thanks in advance.

Answer 1

You can just use product from itertools like,

>>> l
[[1, 2, 3], [4, 5, 6]]
>>> import itertools
>>> list(itertools.product(*l))
[(1, 4), (1, 5), (1, 6), (2, 4), (2, 5), (2, 6), (3, 4), (3, 5), (3, 6)]

Answer 2

What you described is exactly the Cartedian product , and in python it is implemented by the library itertools.
Using itertools.product , you can solve your problem with a single line of code:

import itertools
my_list = [[1, 2, 3], [4, 5, 6]]
list(itertools.product(*my_list))
# output: [(1, 4), (1, 5), (1, 6), (2, 4), (2, 5), (2, 6), (3, 4), (3, 5), (3, 6)]

Or equivantelly, using list comprehension:

list1, list2 = my_list
[(x,y) for x in list1 for y in list2]

A possible product implementation could be the following:

def product(*my_list):
    l_tuple = map(tuple, my_list)
    result = [[]]
    for t in l_tuple:
        result = [a + [b] for a in result for b in t]
    for prod in result:
        yield tuple(prod)