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)