What would be a better implementation of all combinations in lexicographic order of a jagged list?

Question

I was put in a position today in which I needed to enumerate all possible combinations of jagged list. For instance, a naive approach would be:

for a in [1,2,3]:
    for b in [4,5,6,7,8,9]:
        for c in [1,2]:
            yield (a,b,c)

This is functional, but not general in terms of the number of lists that can be used. Here is a more generalized approach:

from numpy import zeros, array, nonzero, max

make_subset = lambda x,y: [x[i][j] for i,j in enumerate(y)]

def combinations(items):
    num_items = [len(i) - 1 for i in items]
    state = zeros(len(items), dtype=int)
    finished = array(num_items, dtype=int)
    yield grab_items(items, state)
    while True:
        if state[-1] != num_items[-1]:
            state[-1] += 1
            yield make_subset(items, state)
        else:
            incrementable = nonzero(state != finished)[0]
            if not len(incrementable):
                raise StopIteration
            rightmost = max(incrementable)
            state[rightmost] += 1
            state[rightmost+1:] = 0
            yield make_subset(items, state)

Any recommendations on a better approach or reasons against the above approach?

Answer 1

This question looks similar:

http://stackoverflow.com/questions/215908/whats-a-good-non-recursive-algorithm-to-calculate-a-cartesian-product

Answer 2

The naive approach can be written more compactly as a generator expression:

((a,b,c) for a in [1,2,3] for b in [4,5,6,7,8,9] for c in [1,2])

The general approach can be written much more simply using a recursive function:

def combinations(*seqs):
  if not seqs: return (item for item in ())
  first, rest = seqs[0], seqs[1:]
  if not rest: return ((item,) for item in first)
  return ((item,) + items for item in first for items in combinations(*rest))

Sample usage:

>>> for pair in combinations('abc', [1,2,3]):
...   print pair
... 
('a', 1)
('a', 2)
('a', 3)
('b', 1)
('b', 2)
('b', 3)
('c', 1)
('c', 2)
('c', 3)