Pythonのジェネレータで遊ぶ + エイトクイーンパズルをジェネレータを使って解いてみる

Pythonでは関数の途中で結果を一旦返して、次に呼び出す時はその続きから処理を再開するような関数を書くことができます。これをジェネレータと呼びます。
ジェネレータがあると例えば、ある関数がリストを返す場合に、一度にリストの全要素を作らなくても、少しずつ呼び出し側に渡してたりできます。つまり、要素数が1,000,000個もあるようなリストを一度に返すとメモリを大量に消費してしまいますが、ジェネレータなら一つ一つ返すので節約ができるわけです。
また、一つ一つの要素を計算するのに膨大な時間がかかるような場合も、ジェネレータなら必要になる度に少しずつ時間を使うだけなので使い方によっては便利ですね。
とは言ったものの、ジェネレータを自分自身で書いたことがなかったので、順列を生成するプログラムを無理矢理ジェネレータで書いてみました。というのも、1〜nの数の順列の組み合わせはn!個なので、n=10でも403,200通りになって、全ての組を生成して一度に返すと酷い目に合うことが目に見えているからです。これはジェネレータの出番ですね！

ジェネレータを書く方法は簡単に言えば、関数でreturnするところをyieldすればいいのです。yieldステートメントを使用すると、関数は呼び出し側に値を渡す時に処理を終了せずに一時停止します。その時、ステート情報が保持されるので、次はその続きから処理を再開できます。ただ、この時何もかもを保持するわけではないので注意が必要です。*1

というわけで、こんなプログラムを書いてみました。(分かりにくくて申し訳ないです)

#!/usr/bin/env python
#-*- coding: utf-8 -*-

# 順列の生成
def generate_perm(n):
    # nは1以上の整数
    assert isinstance(n, int)
    assert n > 0

    # 順列を格納するリスト
    perm = [None for j in range(n)]

    # permの中での位置
    m = 0

    while True:
        x = perm[m] or 0 #None の場合は0をセット

        for y in range(x+1, n+1):
            if y in perm: # すでに使われていたらスキップ
                continue

            perm[m] = y

            if m + 1 < n: # まだ全て選んでいない場合
                m += 1
            else: # 全て選ばれている
                yield perm # ここで一旦結果を返す
                perm[m] = None
                m -= 1
            break
        else:
            # yがnまで行った場合は
            # mを1減らして最初に戻る
            perm[m] = None
            m -= 1

            # 1つめの値がnを越えていたら終了
            if m == -1:
                break

実行してみると

>>> for perm in generate_perm(3):
...     print perm
... 
[1, 2, 3]
[1, 3, 2]
[2, 1, 3]
[2, 3, 1]
[3, 1, 2]
[3, 2, 1]
>>>

ジェネレータはイテレータプロトコルをサポートするオブジェクトなので、nextメソッドを持っています。これをみると、ジェネレータが一度に全てを返すのではないということがよく分かります。

>>> generator = generate_perm(3)
>>> generator
<generator object generate_perm at 0xb7831e3c>
>>> generator.next()
[1, 2, 3]
>>> generator.next()
[1, 3, 2]
>>> generator.next()
[2, 1, 3]
>>> generator.next()
[2, 3, 1]
>>> generator.next()
[3, 1, 2]
>>> generator.next()
[3, 2, 1]
>>> generator.next()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>>

おまけ：エイト・クイーンパズルを解く

せっかくなので、generate_perm関数を使ってエイト・クイーンパズルを解いてみました。エイト・クイーン - Wikipedia
エイト・クイーンパズルを解くアプローチは色々考えられますが、盤面上で2つのクイーンが同じ行もしくは列にくることはないので、この条件から枝刈りができ、このようなクイーンの配置は順列を使って表現することができます。
例えば、4つの順列の一つである[1,3,2,4]や[2,4,1,3]が

     Q * * * | * * Q *
     * * Q * | Q * * *
     * Q * * | * * * Q
     * * * Q | * Q * *

に対応していると考えれば、どのクイーンも同じ行と列にいないのが分かります。
あとは、generate_permが生成する解の候補を一つ一つチェックしていけばいいわけですね。
チェックするための関数とかはここを参考にして書きましたhttp://www.geocities.jp/m_hiroi/light/python03.html#chap08

#!/usr/bin/env python
#-*- coding: utf-8 -*-

def check(perm):
    for y in range(len(perm)):
        if conflict(perm, perm[y], y):
            return False
    return True

def conflict(perm, x, y):
    for y1 in range(y):
        x1 = perm[y1]
        if x1 - y1 == x - y or x1 + y1 == x + y:
            return True
    return False

ではやってみましょう。Wikipediaによると8Queenの場合回転とかを無視すると92通り答えがあるらしいです。

>>> from perm import generate_perm
>>> from eignt_queen import check
>>> c = 0
>>> for perm in generate_perm(8):
...     if check(perm):
...         c += 1
...
>>> print c
92
>>>