Anthony Chung
Anthony Chung

Reputation: 1477

python - turn a combination function into a generator

I have a function that creates all combinations of list items, representing this as a list of lists:

def makeCombos(arr):
  yield (sum([map(list, combinations(arr, i)) for i in range(len(arr) + 1)], []))

Calling makeCombos([1,2,3,4,5]) gives me a generator object, but calling .next() does not give me one combo at a time, it gives me the entire list of combos.

How can I turn this into a generator function that I can call?

Upvotes: 0

Views: 1164

Answers (2)

amiller27
amiller27

Reputation: 506

itertools already has a method for joining iterables together: it's called chain. What you want is something like the following:

def makeCombos(arr):
    return chain.from_iterable(combinations(arr, i) for i in range(len(arr) + 1))

Simple, short, and fairly Pythonic in my opinion.

Upvotes: 1

Pythonista
Pythonista

Reputation: 11635

sum(iterable, []) doesn't create a list of lists. It actually flattens things.

yield (sum(...)) in this line you're just yielding a single item, the flattened list of all combinations.

For Python 2.X sum([(map(list, combinations(arr, i))) ...]) will work, but in Python 3.X map no longer returns a list. Instead it returns a map object. So, if anyones on Python 3.X simply turn this into list(map(.....)) for this to run on 3.X.

I think what you actually want is something like this:

from itertools import combinations
def makeCombos(arr):
     for i in range(len(arr) + 1):
         for combo in map(list, combinations(arr, i)):
            yield combo

#Or call next 

combos = makeCombos([1, 2, 3, 4, 5])
for combo in combos:
     print combo

An alternative from the comment(s) for a one-liner:

Instead of yielding we can return a generator object and cycle through just as we would with the yield.

e.g. -

from itertools import combinations

def makeCombos(arr):

     return (combo for i in range(len(arr) + 1) for combo in map(list, combinations(arr, i)))

combos = makeCombos([1, 2, 3, 4, 5])
....

As for this being "Pythonic" I wouldn't really say so. I actually prefer the nested forloop it is by far more readable.

Although, we can still try to clean it up some more / compact it by doing a few "tricks"

from itertools import combinations as cs #or some other name)

def makeCombos(arr):

    return (c for i in range(len(arr) + 1) for c in map(list, cs(arr, i)))

But, now you've lost all readability and this looks like something you'd see in Perl. (the horror!)

Output:

[]
[1]
[2]
[3]
[4]
[5]
[1, 2]
[1, 3]
[1, 4]
[1, 5]
[2, 3]
[2, 4]
[2, 5]
[3, 4]
[3, 5]
[4, 5]
[1, 2, 3]
[1, 2, 4]
[1, 2, 5]
[1, 3, 4]
[1, 3, 5]
[1, 4, 5]
[2, 3, 4]
[2, 3, 5]
[2, 4, 5]
[3, 4, 5]
[1, 2, 3, 4]
[1, 2, 3, 5]
[1, 2, 4, 5]
[1, 3, 4, 5]
[2, 3, 4, 5]
[1, 2, 3, 4, 5]

Upvotes: 3

Related Questions