Generate combination of subset N from multiple lists

September 7, 2023

I have 4 lists of elements and am trying to generate combinations of elements from those lists. Specifically, I’d like each combination to have 1 element from each list. However, I would only like each combination to have 3 elements in it. I would also like to record which elements are left out.

Example lists:

list1 = ['a', 'b', 'c']
list2 = ['m', 'n', 'o']
list3 = ['x', 'y', 'z']
list4 = ['q', 'r', 's']

Example desired output of just 2 rows but I would like all combinations:

combos = [[['a', 'm', 'x'], 'q'], 
          [['a', 'n', 'r'], 'z'], 
          [['s', 'z', 'o'], 'a']]
    
df = pd.DataFrame(combos, columns = ['combo', 'extra'])

Importantly, the "extras" should be sampled from all lists.

I am currently uncertain if order matters to me or not for the combinations, but if the code is easy to explain for permutations that would also be great. Thank you!

>Solution :

You probably want itertools.product

from itertools import product

ls = [list1, list2, list3, list4]

for *a, b in product(*ls):
    print(a, b)

['a', 'm', 'x'] q
['a', 'm', 'x'] r
['a', 'm', 'x'] s
['a', 'm', 'y'] q
['a', 'm', 'y'] r

Specifically to collect them:

df = pd.DataFrame(
    [[a, b] for *a, b in product(*ls)], 
    columns=['combo', 'extra']
)

Edit

To swap which is the "extra", you’ll want to permute ls:

from itertools import permutations, product


ls = [list1, list2, list3, list4]

for l in permutations(ls):
    for *a, b in product(*l):
        print(a, b)