Pythonでeach_slice

http://blog.livedoor.jp/dankogai/archives/51838970.html

このエントリを見て、PythonRubyのeach_sliceを書くとしたらどうなるだろうと思ってやってみた。

これだと、要素数がnで割り切れない場合は余りの分は返されない。

def each_slice(n, seq):
    next = iter(seq).next
    r = range(n)
    while True:
        yield [next() for i in r]

>>> for x in each_slice(5, range(21)):
>>> print x
[0, 1, 2, 3, 4]
[5, 6, 7, 8, 9]
[10, 11, 12, 13, 14]
[15, 16, 17, 18, 19]

これなら余りも返ってくるけど、StopIterationが伝播されないからシーケンスの終了を自分でチェックする必要がある。

def each_slice(n, seq):
    next = iter(seq).next
    r = range(n)
    item = list(next() for i in r)
    while item:
        yield item
        item = list(next() for i in r)

>>> for x in each_slice(5, range(21)):
>>> print x
[0, 1, 2, 3, 4]
[5, 6, 7, 8, 9]
[10, 11, 12, 13, 14]
[15, 16, 17, 18, 19]
[20]

groupbyを使ったパターン。無駄な計算が入るし効率はよくないよね。

from itertools import groupby
def each_slice(n, seq):
    for k, group in groupby(enumerate(seq), lambda x:x[0]//n):
        yield [x[0] for x in group]

>>> for x in each_slice(5, range(21)):
>>> print x
[0, 1, 2, 3, 4]
[5, 6, 7, 8, 9]
[10, 11, 12, 13, 14]
[15, 16, 17, 18, 19]
[20]

全要素数がnで割り切れない場合もちゃんと最後まで要素nのリストを返す方が使いやすいから、そうなるとこうかな。

from itertools import chain
def each_slice(n, seq):
    next = iter(chain(seq, [None]*(n-1))).next
    r = range(n)
    while True:
        yield [next() for i in r]

>>> for x in each_slice(5, range(21)):
>>> print x
[0, 1, 2, 3, 4]
[5, 6, 7, 8, 9]
[10, 11, 12, 13, 14]
[15, 16, 17, 18, 19]
[20, None, None, None, None]

結局、itertoolsのドキュメントに載ってるgrouperが一番すっきりしてますね。

from itertools import izip_longest
def each_slice(n, seq):
    args = [iter(seq)] * n
    return izip_longest(fillvalue=None, *args)

(参考) http://stackoverflow.com/questions/3833589/python-equivalent-of-rubys-each-slicecount