KevinShaffer
KevinShaffer

Reputation: 858

Python Recursive Generator

I'm trying to make a recursive generator, but I'm clearly missing something. Basically I have a list of positive and negative values for specific dimensions of a histogram, and want to get every combination of positive and negative. Ex:

input = [[-1,2],[-3,4],[-5,6]]

And the result I need is (although I want this as a generator instead of a list, but you get the idea):

output = [[-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 managed to get it done, but only by appending items to a global list inside of my recursion, but I really was hoping I could just do this with a generator and iterate over my function since that would make it so much cleaner IMO.

Here's what I have now:

def combs(idx,array,current=[]):
    for item in array[idx]:
        current_copy = current[:]
        current_copy.append(item)
        if idx + 1 < len(array):
            combs(idx+1,array,current_copy)
        else:
            print current_copy
            # yield current_copy

This will print out exactly what I want, but when I try to change the print to a yield and loop over the function it doesn't work. (e.g.)

for item in combs(0,input,[]):
    print item

Upvotes: 0

Views: 1252

Answers (2)

glibdud
glibdud

Reputation: 7840

Another option for this task is the itertools.product function:

>>> from itertools import product
>>> input = [[-1,2],[-3,4],[-5,6]]
>>> list(product(*input))
[(-1, -3, -5), (-1, -3, 6), (-1, 4, -5), (-1, 4, 6), (2, -3, -5), (2, -3, 6), (2, 4, -5), (2, 4, 6)]

Upvotes: 2

Chris Lawlor
Chris Lawlor

Reputation: 48902

If you are using Python 3.3 or newer, you want the yield from statement:

def combs(idx,array,current=[]):
    for item in array[idx]:
        current_copy = current[:]
        current_copy.append(item)
        if idx + 1 < len(array):
            yield from combs(idx+1,array,current_copy)
        else:
            yield current_copy

On Python 2.7, you can expand the yield from into a loop:

def combs(idx,array,current=[]):
    for item in array[idx]:
        current_copy = current[:]
        current_copy.append(item)
        if idx + 1 < len(array):
            for i in combs(idx+1,array,current_copy):
                yield i
        else:
            yield current_copy

Upvotes: 2

Related Questions