イテラブル(list, tuple 等)から数個ずつ取り出したい(Python)

こんにちは。今回はイテラブルから特定の個数ずつ要素を取り出していく方法について語ろうと思います。「イテラブルとはなんぞや?」という方もいらっしゃると思いますが、ざっくりいうと for 文を使って要素を取り出せるようなオブジェクトです。

# こういうのが可能なやつ
for x in some_iterable:
    print(x)

例えばリストとかタプルとかがイテラブルです。まずはリストを数個ずつ切り出す方法から初めて、その後一般のイテラブルについて議論します。

リスト編

スライスを使う方法

たとえば以下のようなリストがあったとします:

a = [0, 1, 2, 3, 4, 5]

これを2個ずつに分離する最も簡単な方法はスライスでしょう:

print(a[0:2])  # [0, 1]
print(a[2:4])  # [2, 3]
print(a[4:6])  # [4, 5]

for ループで書くならこんな感じ:

for i in range(0, len(a), 2):
    print(a[i:i+2])  # [0, 1], [2, 3], [4, 5]

これらを集めたリストを内包表記で作るとこんな感じ:

b = [a[i:i+2] for i in range(0, len(a), 2)]
for x in b:
    print(x)  # [0, 1], [2, 3], [4, 5]

簡単ですね。これを関数にしてしまいましょう。

def divide_into_pieces_0(x: list, size: int) -> list:
    return [x[i:i + size] for i in range(0, len(x), size)]


a = [0, 1, 2, 3, 4, 5]
for piece in divide_into_pieces_0(a, 2):
    print(piece)  # [0, 1], [2, 3], [4, 5]

スライスと len() を使えればよいので、リストの代わりにタプルを受け取ってもちゃんと動作します。

インデックスを使わない方法

「各ピースの先頭のインデックスを取り出し、欲しい個数を切り出す」というのが先ほどの手法でした。しかし「わざわざ先頭のインデックス探すのダサくない・・・?」という気もします。切り出す位置のインデックスを一旦使わない方法として、こういうのも可能です:

a = [0, 1, 2, 3, 4, 5]
print(a[:2])  # [0, 1]
a = a[2:]  # [2, 3, 4, 5]
print(a[:2])  # [2, 3]
a = a[2:]  # [4, 5]
print(a[:2])  # [4, 5]

個数を表す「2」だけで操作できるようになりました。全部出し切るまで行うことに注意してループで書くとこんな感じです:

while a:  # a に要素がある限り
    print(a[:2])
    a = a[2:]  # 最後は a == [] となって終了

これも関数化しちゃいましょう。

def divide_into_pieces_1(x: list, size: int) -> list:
    pieces = []
    while x:
        pieces.append(x[:size])
        x = x[size:]  # オブジェクト生成するので引数への影響なし
    return pieces


a = [0, 1, 2, 3, 4, 5]
for piece in divide_into_pieces_1(a, 2):
    print(piece)  # [0, 1], [2, 3], [4, 5]

append を避けてジェネレータ式を使ってもよいでしょう:

from typing import Generator

def divide_into_pieces_2(x: list, size: int) -> Generator[list, None, None]:
    while x:
        yield x[:size]  # for で取り出す要素になる
        x = x[size:]


a = [0, 1, 2, 3, 4, 5]
for piece in divide_into_pieces_2(a, 2):
    print(piece)  # [0, 1], [2, 3], [4, 5]

スライスを避ける方法

for 文でどんどん取り出して貯めていき、特定の個数が溜まったら出荷する、という発想もできます:

a = [0, 1, 2, 3, 4, 5]
piece = []
for i, x in enumerate(a):
    piece.append(x)
    if i % 2 == 1:
        print(piece)  # [0, 1], [2, 3], [4, 5]
        piece = []  # 出荷が終わったので空にする

この実装はひとつ問題があり、a の要素数が奇数個だと最後の一個が出荷されず闇に葬られていまいます。そこで少しだけ工夫して葬られかけたものを救っておきましょう:

a = [0, 1, 2, 3, 4, 5, 6]
piece = []
for i, x in enumerate(a):
    piece.append(x)
    if i % 2 == 1:
        print(piece)  # [0, 1], [2, 3], [4, 5]
        piece = []  # 出荷が終わったので空にする
if piece:  # 闇に葬られそうなら助ける
    print(piece)  # [6]

同じく関数にするとこんな感じです。

def divide_into_pieces_3(x: list, size: int) -> list:
    pieces = []
    piece = []
    for i, e in enumerate(x):
        piece.append(e)
        if i % size == (size - 1):
            pieces.append(piece)
            piece = []  # 新しい空のリストにする
    if piece:
        pieces.append(piece)
    return pieces


a = [0, 1, 2, 3, 4, 5]
for piece in divide_into_pieces_3(a, 2):
    print(piece)  # [0, 1], [2, 3], [4, 5]

リストのリストではなく、リストを生み出すジェネレータを作るならこんな感じでしょうか:

from typing import Generator

def divide_into_pieces_4(x: list, size: int) -> Generator[list, None, None]:
    piece = []
    for i, e in enumerate(x):
        piece.append(e)
        if i % size == (size - 1):
            yield piece
            piece = []
    if piece:
        yield piece


a = [0, 1, 2, 3, 4, 5]
for piece in divide_into_pieces_4(a, 2):
    print(piece)  # [0, 1], [2, 3], [4, 5]

一般のイテラブル編

この記事の冒頭で「イテラブルとは for 文が使えるオブジェクトだ」という風に説明しました。より正確には __iter__ メソッドが定義されていて、iter(obj) でイテレーターを生成できるようなものがイテラブルです(多分)。イテレーターについては過去に紹介したことがあるので↓よければ読んでみてください。
smooth-pudding.hatenablog.com
この記事ではイテラブルがどのように挙動するかについては既知とします。
なおイテラブルを引数にもらって最初にリストに変換(list のコンストラクタに入れるだけで作れます)して、リスト編のテクニックをそのまま使うことも可能です。重複を避けるため、以下では「引数をもらったら速攻でリストにする」という方法を縛ります。これはめちゃくちゃ長いイテレータが来たときに困らないようにするためでもあります。

溜まったら yield する作戦

リスト編の最後に紹介した方法では、もはやリストのインデックスを全く使いませんでした。なので一般のイテラブルに適用可能です。

from typing import Generator, Iterable

def divide_into_pieces_5(x: Iterable, size: int) -> Generator[list, None, None]:
    piece = []
    for i, e in enumerate(x):
        piece.append(e)
        if i % size == (size - 1):
            yield piece
            piece = []
    if piece:
        yield piece


a = [0, 1, 2, 3, 4, 5]
for piece in divide_into_pieces_5(a, 2):
    print(piece)  # [0, 1], [2, 3], [4, 5]

enumerate を使わない方法

上記の方法はインデックスの代わりに enumerate を使ってインデックスを生成していました。リスト編ではインデックスを使わずに取り出す方法を紹介したので、一般のイテラブルでも enumerate を縛ってみましょう。例えばこんなやり方があります:

from typing import Generator, Iterable

def divide_into_pieces_6(x: Iterable, size: int) -> Generator[list, None, None]:
    piece = []
    for e in x:
        piece.append(e)
        if len(piece) == size:
            yield piece
            piece = []
    if piece:
        yield piece


a = [0, 1, 2, 3, 4, 5]
for piece in divide_into_pieces_6(a, 2):
    print(piece)  # [0, 1], [2, 3], [4, 5]

溜まった判定をリストの長さで行うという発想です。

特定の個数ずつで切れるイテレータを作る方法

ここまでは比較的簡単な構文で生成してきました。ここからはイテレータオブジェクトを定義するなどしてより内部構造に足を踏み入れていきます。「特定の個数ずつで切れるイテレータ」とは、こんな感じで挙動するイテレータのイメージです:

class PieceIterator:  # なんかいい感じに定義したイテレータクラス
    pass


a = [0, 1, 2, 3, 4, 5]
it = PieceIterator(a, 2)  # いいかんじイテレータオブジェクトをつくる
print([x for x in it])  # [0, 1]
print([x for x in it])  # [2, 3]
print([x for x in it])  # [4, 5]
print([x for x in it])  # []  # 使い切ったら空になる

要は初期化の際に指定した回数ごとに raise StopIteration するクラスを定義すればOKです。以下が実装の一例です:

from typing import Iterable

class PieceIterator:
    def __init__(self, iterable: Iterable, size: int):
        self.iterator = iter(iterable)  # イテレータオブジェクトを保存しておく
        self.size = size
        self.counter = 0  # 個数を数えるためのカウンタを用意

    def __iter__(self):
        return self  # 「イテレータ?俺だよ」

    def __next__(self):
        if self.counter == self.size:
            self.counter = 0  # raise して処理が中止するまえにカウンタをリセット
            raise StopIteration
        x = next(self.iterator)  # 要素がなければ raise StopIteration
        self.counter += 1  # 要素があればカウンタを増やして...
        return x  # 要素を返す。

このクラスを使って特定の個数ずつ取り出すやつをやってみましょう。

from typing import Generator, Iterable

def divide_into_pieces_7(x: Iterable, size: int) -> Generator[list, None, None]:
    it = PieceIterator(x, size)
    while y := list(it):  # 空になる直前まで続ける
        yield y

a = [0, 1, 2, 3, 4, 5]
for piece in divide_into_pieces_7(a, 2):
    print(piece)  # [0, 1], [2, 3], [4, 5]

y := list(it) の構文は Python3.9 で新たに追加されました。Python3.8までであれば例えば以下のような書き方でokです。

def divide_into_pieces_7(x: Iterable, size: int) -> Generator[list, None, None]:
    it = PieceIterator(x, size)
    while True:
        y = list(it)
        if not y:
            break
        yield y

リストではなくイテレータの断片をリターンしたい→実は不可能

ここまでの方法ではイテラブルをもらってリストを返す実装をしてきました。ここまで来たら「リターンするものもイテレータにできないか?」と考えるのは自然です。しかしよく考えると実はイテレータから細切れのイテレータを取り出すのは不可能で、無理やり作っても意図した挙動を実装できないことが分かります。
例えばこのような実装を考えます:

from typing import Generator, Iterable

def divide_into_pieces_8(x: Iterable, size: int) -> Generator[list, None, None]:
    it = PieceIterator(x, size)
    while True:  # 終了条件??
        yield (e for e in it)

先ほどと違って、yield の終了条件を与えられていないことが分かります。終了条件を与えるには「元のイテレータを使い果たしたら終わり」をうまく表現する必要がありますが、そのためには元のイテレータから発された StopIteration を検出する必要があります。ところでこの StopIteration が raise されるのは文字通りイテレータが使い果たされたタイミングですが、上記のようにイテレータを作って返すときは「実際には要素を取り出していない」ため、どうあがいても使い果たしたかどうかを判定できません(使い果たしていないため)。さらに悪いことに、イテレータそのものは「自分がどこのピースを取り出したものか」という情報を持っていません。つまり上記の関数でイテレータを取り出しても、「出した順に手前からの断片に対応している」というものにはなっていません:

a = [0, 1, 2, 3, 4, 5]
it = divide_into_pieces_8(a, 2)
x = next(it)
y = next(it)
z = next(it)
print(list(y))  # [0, 1]
print(list(z))  # [2, 3]
print(list(x))  # [4, 5]
# どの順序で呼び出してもこの順番に取り出される

x, y, z は「PieceIterator にしたがって手前から取り出す」という役割のみしか知らないため、どこから取り出すのか制御できません。以上の問題を解決するには元のイテラブルに「長さ」や「位置」といった情報を付け加える必要があり、もはやリストやタプルのような特別なイテラブルでなければ対処できないことが分かります*1。なおリストやタプルから細切れのイテレータを取り出すことは可能です(リスト編で作った細切れリストたちに iter() を作用させるだけ)。

まとめ

以上、イテラブルを断片化して生成する方法をまとめてみました。特定の要素数ごとに区切りたいときはたまに発生すると思うので、そんなときに思い出して参考にしてもらえると嬉しいです。

*1:イテレータをコピーしまくって適当に要素を捨てて開始位置を調整し、また「分身で長さをはかる」ことで空かどうか判定するなど、完全に不可能ではないと思いますが、そこまで頑張る意味がどこまであるかは不明です