0

I have a rather complex problem (or so it seems to me). I have a dictionary with key-ed lists of elements and list of the keys of that dictionary

numbers = { 
            'group1': [1, 2],
            'group2': [3, 4],
            'group3': [5, 6]
          }
order   = [ 'group1', 'group2', 'group3' ]

I want to generate combinations by using one number per group by following the order given. So with the order [ 'group1', 'group2', 'group3' ] I expect to get

1, 3, 5
1, 3, 6
1, 4, 5
1, 4, 6
2, 3, 5
2, 3, 6
2, 4, 5
2, 4, 6

I have a solution but it's not generic enough and doesn't scale:

group1=order[0]
group2=order[1]
group3=order[2]
for n in numbers[group1]:
    for n2 in numbers[group2]:
        for n3 in numbers[group3]:
            print(n, n2, n3)

Any smarter solution that will work for an arbitrary number of groups?

MFerguson
  • 1,739
  • 9
  • 17
  • 30
Pithikos
  • 18,827
  • 15
  • 113
  • 136

1 Answers1

1

You can use a combination of itertools.product, a list comprehension and tuple unpacking:

import itertools

for ns in itertools.product(*[numbers[g] for g in order]):
    print(*ns) # or 'print ns' for 2.x

This gives me the result you're looking for:

1 3 5
1 3 6
1 4 5
1 4 6
2 3 5
2 3 6
2 4 5
2 4 6
jonrsharpe
  • 115,751
  • 26
  • 228
  • 437