PythonでRubyのArray#each_sliceみたいなもの

Ruby の配列(Array)には each_slice というメソッドがある。配列のスライスを順に返してくれるイテレータだ。こんなふう。

irb(main):001:0> [0,1,2,3,4,5].each_slice(2){|x| p x}
[0, 1]
[2, 3]
[4, 5]
=> nil

昨日、CodeEval の問題を解くのに、この Array#each_slice が使えると思っていつもの Python じゃなくて Ruby で解いたんだけど、Python には Ruby の Array#each_slice に相当するものがないみたい。

で、今日になってジェネレータを使えばいいんじゃないかと気が付いたので、書いてみた。

def each_slice(lis, n):
    s = 0
    while s < len(lis):
        yield lis[s:s+n]
        s += n

for x in each_slice([0,1,2,3,4,5], 2):
    print x
^o^ >python each_slice.py
[0, 1]
[2, 3]
[4, 5]

上の例では、要素数がちょうど割り切れる数だったけど、半端が出る場合はどうだろう。Ruby だとこうなる。

irb(main):002:0> [0,1,2,3,4,5,6].each_slice(2){|x| p x}
[0, 1]
[2, 3]
[4, 5]
[6]
=> nil

Python で書いたスクリプトだと:

^o^ > python each_slice.py
[0, 1]
[2, 3]
[4, 5]
[6]

IndexError になるかと思ったけど、期待通りに動いた。