無限オブ無限

ナナメに列挙するやり方で書いてみた。

from itertools import count

def flatten(L):
    D = {}
    cnt = 1
    while True:
        for i, gen in enumerate(L):
            g = D.setdefault(i, iter(gen))
            yield g.next()
            if i+1 == cnt:
                break
        cnt += 1

L = [count(0), count(10), count(20), count(30)]
g = flatten(L)
for i in range(10):
    print g.next(),

# 出力結果
0 1 10 2 11 20 3 12 21 30

ただ不可思議だったのが、無限リストの無限リスト。挙動が何かおかしい。Pythonは無限リストの無限リストを表せない?

from itertools import count

def cnt_gen():
    i = 0
    while True:
        yield count(i)
        i += 10

L = cnt_gen()
for i in range(3):
    gen = L.next()
    print id(gen), gen

# 出力結果
11632880 count(0)
11632896 count(10)
11632880 count(20)


追記:

イテレータが空になったときエラーで落ちていたので少し修正。あまりテストしていないが、これで有限でも無限でも問題ないはず。

def flatten(L):
    D, cnt = {}, 1
    while True:
        for i, gen in enumerate(L):
            g = D.setdefault(i, iter(gen))
            try:
                yield g.next()
            except StopIteration:
                pass
            if i+1 == cnt:
                break
        cnt += 1