Какова была бы лучшая реализация всех комбинаций в лексикографическом порядке зубчатого списка?

Question

Сегодня меня поставили в положение, в котором мне нужно было перечислить все возможные комбинации зубчатого списка. Например, наивный подход будет:

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

Это функционально, но не является общим с точки зрения количества списков, которые можно использовать. Вот более обобщенный подход:

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)

Какие-либо рекомендации по лучшему подходу или причины против вышеупомянутого подхода?

Answer 1

Наивный подход можно записать более компактно в виде выражения генератора:

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

Общий подход можно написать гораздо проще, используя рекурсивную функцию:

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))

Пример использования:

>>> 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)