ジェネレータが返した値が消える理由

はじめに

Scheme で継続の練習としてバックトラック法を使うコードを書いていたんですが、そもそもバックトラック法を理解していないなと思ったので Python で試したらいろいろわかったという記事です。

アルゴリズム

バックトラック法は全探索の一種で、探索が失敗すると一つ手順を戻して次の探索をすることを繰り返す手法。特に木構造についてこの手法を使う場合は深さ優先探索 (DFS) と呼ばれる。

例題としてピタゴラス数を使う。
サンプルコードは以下。

from typing import Callable, Iterator, Sequence, TypeVar

T = TypeVar("T")

def pythagorean_backtrack(
    pred: Callable[..., bool],
    domains: Sequence[Sequence[T]]
) -> Iterator[tuple[T, ...]]:

    length: int = len(domains)

    def dfs(idx: int, chosen: list[T]) -> Iterator[tuple[T, ...]]:
        if idx == length:
            if pred(*chosen):
                yield tuple(chosen)
            return

        for value in domains[idx]:
            chosen.append(value)
            yield from dfs(idx + 1, chosen)
            chosen.pop()

    yield from dfs(0, [])

if __name__ == "__main__":
    gen = pythagorean_backtrack(
        lambda x, y, z: x**2 + y**2 == z**2,
        [
            [i for i in range(1, 40)],
            [i for i in range(1, 40)],
            [i for i in range(1, 40)],
        ],
    )
    for ret in gen:
        print(ret)

lambda x, y, z: x**2 + y**2 == z**2 では枝刈りができない、記事を通して枝刈りの話を全くしていないが、pred を引数で渡すようにしているので自由に枝刈りされたい。

コードをいろいろ変更しているうちに

    def dfs(idx: int, chosen: list[T]) -> Iterator[tuple[T, ...]]:
        if idx == length:
            if pred(*chosen):
                yield tuple(chosen)
            return

を

    def dfs(idx: int, chosen: list[T]) -> Iterator[list[T, ...]]:
        if idx == length:
            if pred(*chosen):
                yield chosen
            return

に書きかえると

for ret in gen:
    print(ret)
# [3, 4, 5]
# [4, 3, 5]
# ... 
 
print(list(gen))
# [[], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], []]

になることに気がついた。
両方が動かないのならともかく片方だけ？と思った。

for ループと list() によるジェネレータの消費は結果としては同じデータを生成できるからだ。

(今回はたまたま文字列とリスト型の違いがあるが、やろうとすれば同じ型も生成できる)

なぜか

実はこれには Python における引数の評価戦略が関係していることがわかった。

先の２つの場合は、print() による評価のタイミングが違う。

• for ループの場合は要素が排出されると即 print() によって評価される
• list() の場合は全ての要素が吐き出されたあと print() に評価される

Python は引数の値呼びをする言語ではないため、list(gen) の中身は実はオブジェクトへの参照だ。ジェネレータの内部では yield で返されたあとも同じオブジェクトを使いまわしており、pop() された結果、配列の中身も闇に消えてしまったのだ。

それが分かれば理由は明確になる。次のコードを走らせて見ればなぜ tuple() が必要かわかるだろう。

pythagorean2.py

def pythagorean_backtrack(
    pred: Callable[..., bool],
    domains: Sequence[Sequence[T]]
) -> Iterator[tuple[T, ...]]:

    length: int = len(domains)

    def dfs(idx: int, chosen: list[T]) -> Iterator[tuple[T, ...]]:
        if idx == length:
            if pred(*chosen):
                print("chosen", id(chosen))
                print("tuple(chosen)", id(tuple(chosen)))
                yield tuple(chosen)
            return

        for value in domains[idx]:
            chosen.append(value)
            yield from dfs(idx + 1, chosen)
            chosen.pop()

    yield from dfs(0, [])


if __name__ == "__main__":
    gen = pythagorean_backtrack(
        lambda x, y, z: x**2 + y**2 == z**2,
        [
            [i for i in range(1, 40)],
            [i for i in range(1, 40)],
            [i for i in range(1, 40)],
        ],
    )
    list(gen)

見やすさのために、出力を sort | uniq -c に通して集計する。

result.sh

$ uv run  ./pythagorean2.py | sort | uniq -c
     30 chosen 132498275353984
      1 tuple(chosen) 132498275409072
      1 tuple(chosen) 132498275410192
      1 tuple(chosen) 132498275411152
      1 tuple(chosen) 132498275411232
      1 tuple(chosen) 132498275411312
      1 tuple(chosen) 132498275411392
      1 tuple(chosen) 132498275411472
      1 tuple(chosen) 132498275411552
      1 tuple(chosen) 132498275411632
      1 tuple(chosen) 132498275411712
      1 tuple(chosen) 132498275411792
      1 tuple(chosen) 132498275411872
      1 tuple(chosen) 132498275411952
      1 tuple(chosen) 132498275412032
      1 tuple(chosen) 132498276373648
      1 tuple(chosen) 132498276728752
      1 tuple(chosen) 132498276923520
      1 tuple(chosen) 132498276924720
      1 tuple(chosen) 132498277026688
      1 tuple(chosen) 132498277030848
      1 tuple(chosen) 132498277031168
      1 tuple(chosen) 132498277031248
      1 tuple(chosen) 132498277031328
      1 tuple(chosen) 132498277031488
      1 tuple(chosen) 132498277031888
      1 tuple(chosen) 132498277034048
      1 tuple(chosen) 132498277034608
      1 tuple(chosen) 132498277038688
      1 tuple(chosen) 132498277039248
      1 tuple(chosen) 132498277039648

このように配列の中身が同じ id であること、tuple() が新しいオブジェクトを作り id が変わることがわかる。

終わりに

中身を入れたり出したりする処理であるからリスト型の方が向いているだろうと思って使っているが、実際はタプル型だけでも書ける。

毎周、オブジェクトを作るのは重そうな気がするがサンプルコードの範囲だと実行速度はあまり変わらない。

また返り値の型が変わるのを許すのなら tuple() の代わりに list() を使っても構わない。手元の環境では、むしろこちらの方が少し速かった。

def pythagorean_backtrack_tuple(
    pred: Callable[..., bool],
    domains: Sequence[Sequence[T]]
) -> Iterator[tuple[T, ...]]:
    """
    domains で与えられた各候補集合から値を1つずつ選び、
    その組み合わせのうち pred が True を返すものだけを列挙するジェネレータ。

    max_length: ここで止める
    """

    max_length: int = len(domains)

    def dfs(idx: int, chosen: tuple[T, ...]) -> Iterator[tuple[T, ...]]:
        """
        再帰的に探索を行う内部関数。

        idx   : 現在選択中の変数のインデックス
        chosen: これまでに選んだ値のタプル
        """

        # すべての変数を選び終えた場合
        if idx == max_length:
            # 条件を満たしていれば結果として返す
            if pred(*chosen):
                yield chosen
            return

        # domains[idx] にある候補を1つずつ試す
        for value in domains[idx]:
            yield from dfs(idx + 1, (*chosen, value))

    # 初期状態は空タプルからスタート
    yield from dfs(0, ())


if __name__ == "__main__":
    gen = pythagorean_backtrack_tuple(
        lambda x, y, z: x**2 + y**2 == z**2,
        [
            [i for i in range(1, 40)],
            [i for i in range(1, 40)],
            [i for i in range(1, 40)],
        ],
    )
    print(list(gen))

Link

• バックトラッキング - Wikipedia
• 深さ優先探索 - Wikipedia
• 幅優先探索 - Wikipedia
• 評価戦略 - Wikipedia