組み合わせを生成するジェネレータ

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

のような入力から、

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

と、全ての組み合わせを生成するジェネレータがPythonで必要になったので書いてみた
(本当はタプルで返したほうがよい)

ちょっと書くのに時間がかかったので、紛失した時用にメモ

from __future__ import print_function

def combinations( nested ,fix=[]):

    if nested==[]:

        yield fix

    else:

        car=nested[0]

        cdr=nested[1:]

        if not isinstance(car,(list,tuple)):

            car=[car]

        for i in car:

            fix2=fix+[i]

            for j in combinations(cdr,fix2):

                yield j

テストコード

if __name__=="__main__":

    testdata=[]

    for x in [1,2,3]:

        for y in [4,5,6]:

            for z in [7,8,9]:

                testdata.append( [x,y,z,10])

    arg=[ [1,2,3],[4,5,6],[7,8,9],10]

    ret=list(combinations(arg))

    assert ret==testdata

    if True:

        print( arg)

        for r in ret:

            print (r)

    print ("done")

実行結果

[[1, 2, 3], [4, 5, 6], [7, 8, 9], 10]

[1, 4, 7, 10]

[1, 4, 8, 10]

[1, 4, 9, 10]

[1, 5, 7, 10]
〜中略〜
[3, 6, 8, 10]

[3, 6, 9, 10]

done