itertools and selection from partitioned sets

Question

I have a list of lists, which we can think of as 3 urns with 3 balls.

mylist= [
            ['HIGH_1', 'MED_1', 'LOW_1'],
            ['HIGH_2', 'MED_2', 'LOW_2'],
            ['HIGH_3', 'MED_3', 'LOW_3']
        ]

I'm only allowed to select one ball from each urn. I would like to get all combinations of 1 ball, 2 balls, and 3 balls. Answer in a list of lists.

In order to achieve this I place a fake element (None) in each urn

mylist= [
            [None, 'HIGH_1', 'MED_1', 'LOW_1'],
            [None, 'HIGH_2', 'MED_2', 'LOW_2'],
            [None, 'HIGH_3', 'MED_3', 'LOW_3']
           ]

Now the following use of itertools and filter gets me the desired solution.

combos = []
for l in itertools.product(*mylist):
    combos.append(filter(lambda a: a is not None, l))

combos.remove(())   # remove the empty element
print [list(elem) for elem in combos]


[['HIGH_3'], ['MED_3'], ['LOW_3'], ... , ['MED_2', 'HIGH_3'], ['MED_2', 'MED_3'], ['MED_1', 'HIGH_3'], ..., ['LOW_1', 'LOW_2', 'MED_3'], ['LOW_1', 'LOW_2', 'LOW_3']]

Which produces 63 elements.

9 (single elements) + 27 (two elements) + 27 (three elements) = 63

Adding in that fake None and post filtering doesn't feel like the best way to do this.

Is there a way to avoid those two seemingly redundant steps?

Answer 1

You can do it with the following generator

def combos(data):
    for i in xrange(1, len(data) + 1):
        for item in itertools.combinations(data, i):
            for j in itertools.product(*item):
                yield j

Probably the most pythonic solution I can think of.