관리-도구

편집 파일: IterPermute.py

from functools import reduce

def iterpermute(*args):
    args = [iter(a) for a in args]
    la = len(args)
    stopped = [0] * la
    lens = [0] * la
    bufs = [[] for i in range(la)]
    nexts = [None] * la
    n = 0
    while 1:
        anynew = 0
        for i in range(la):
            if stopped[i]:
                nextval = bufs[i][n % lens[i]]
            else:
                try:
                    nextval = next(args[i])
                except StopIteration:
                    if lens[i] == 0:
                        # raise ValueError, 'The iterator passed in arg %d did not return any item'%i
                        return
                    stopped[i] = 1
                    nextval = bufs[i][n % lens[i]]
                else:
                    anynew = 1
                    bufs[i].append(nextval)
                    lens[i] += 1
            nexts[i] = nextval
        if anynew:
            n += 1
            yield tuple(nexts)
        else:
            break

wanted = reduce(lambda x, y: x*y, lens, 1)
    if n >= wanted:
        assert n == wanted
        return
    ixs = list(enumerate(lens))
    ixs.sort(key=lambda x: x[1])
    ixs = [ix for (ix, ln) in ixs]
    jxs = [0] * la
    seen = dict([(tuple([j % lens[i] for i in ixs]), 1)
                 for j in range(n)])

while n < wanted:
        t = tuple([jxs[i] for i in ixs])
        if t not in seen:
            yield tuple([bufs[i][jxs[i]] for i in range(la)])
            n += 1

for i in ixs:
            j = jxs[i]
            j = (j + 1) % lens[i]
            jxs[i] = j
            if j != 0:
                break

def test_iterpermute():
    import itertools
    repeat = itertools.repeat
    assert list(iterpermute()) == [()]
    assert list(iterpermute(repeat(1, 2))) == [(1,), (1,)]
    assert list(iterpermute(repeat(1, 1), repeat(2, 1))) == [(1, 2)]
    assert list(iterpermute(list(range(0, 2)), list(range(2, 3)))) == [
        (0, 2), (1, 2)]
    assert list(iterpermute(list(range(0, 2)), list(range(2, 4)))) == [
        (0, 2), (1, 3), (1, 2), (0, 3)]
    print(list(iterpermute(list(range(0, 2)), list(range(0, 3)))))
    print(list(iterpermute(list(range(0, 3)), list(range(0, 2)))))

if __name__ == '__main__':
    test_iterpermute()