【Python】イテレータ・ジェネレータからN件ずつ取得する方法

はじめに

みずほリサーチ＆テクノロジーズの@fujineです。

リストや辞書のように、要素を１つずつ取得可能なオブジェクトのことを、Pythonではイテラブルオブジェクトと呼びます。

for文にイテラブルオブジェクトを使用すると、以下のように1件ずつ取得できます。

[i for i in range(5)]  # 0, 1, 2, 3, 4

これを、例えばN=2とした場合に

(0, 1), (2, 3), (4,)

のようにN件ずつ取得する方法を本記事で模索していきます。

検証環境

Python3.9で検証していきます。
使用するデータは、イテラブルオブジェクトであるiterableと、同じ値を生成するジェネレータのgeneratorの2種類です。
一度に取得する要素数はN=2とします。

iterable = [0, 1, 2, 3, 4]
generator = (i for i in iterable)
N = 2

①スライスを使う

インデックスで取得範囲を指定する方法です。最も直感的な実装であり、他のプログラミング言語でもよく用いられます。

def niter_slice(iterable, n):
    return [iterable[i:i+n] for i in range(0, len(iterable), n)]

niter_slice(iterable, N)
# [[0, 1], [2, 3], [4]]

ただし、ジェネレータにこの方法は使えません。ジェネレータは通常__len__()をサポートしていないため、以下のようにエラーとなってしまいます。

niter_slice(generator, N)
# TypeError: object of type 'generator' has no len()

②whileループを使う

whileループを使ったパターンです。①よりもコード量は増えてしまいましたが、ジェネレータにも適用できる点がメリットです。

def niter_while(iterable, n):
    res = []
    it = iter(iterable)
    is_stop = False
    while not is_stop:
        chunk = []
        for _ in range(n):
            try:
                chunk.append(next(it))
            except StopIteration:
                is_stop = True
        if chunk:
            yield chunk

list(niter_while(iterable, N))
# [[0, 1], [2, 3], [4]]

list(niter_while(generator, N))
#[[0, 1], [2, 3], [4]]

③more-itertools.chunkedを使う

more-itertoolsを使うと、標準ライブラリのitertoolsよりも高度で多様なイテラブル関数を利用できます。

以下コマンドでインストールします。

pip install more-itertools

more_itertools.chunkedでN件ずつ取得します。ジェネレータにも適用でき、コードもかなりシンプルです！

import more_itertools

list(more_itertools.chunked(iterable, N))
# [[0, 1], [2, 3], [4]]

list(more_itertools.chunked(generator, N))
# [[0, 1], [2, 3], [4]]

④Numpy.array_splitを使う

取得後のデータをNumpyの配列にするのであれば、Numpy.array_splitが便利です。

以下コマンドでインストールします。

pip install numpy

array_splitの第2引数には、分割サイズではなく分割数を指定します。分割結果はnumpy.ndarray型として取得されます。

import numpy as np

np.array_split(iterable, len(iterable) // N + 1)
# [array([0, 1]), array([2, 3]), array([4])]

残念ながらジェネレータには適用できません。

np.array_split(generator, 3)
# TypeError: object of type 'generator' has no len()

まとめ

色々と模索した結果、汎用性や使い勝手の点ではmore_itertoolsがとても優秀でした。N件取得以外にも様々なイテレーション機能が充実しているようですので、改めて記事化したいと思います。