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!
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']
)
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)