ジェネレーターによる与えられた条件をみたす数列の全列挙

概要

与えられた条件をみたす数列を全列挙したいとき、ジェネレーターを使うと簡単に書けて便利。

ジェネレーターとは

イテレータを生成してくれる関数。return の代わりに yield を使って値を返す。 詳しくは公式ドキュメントを参照されたい。呼ばれる度に次の値を返す。
例えば、

def generate():
    yield 0
    yield 1
    yield 2

で定義すると、

for a in generate():
    print(a)

のように使うことができる。

順列の列挙

$1$ から $n$ までの自然数から $ r $ 個を選んで並べてできる数列をすべて列挙したいとき、Python なら itertools ライブラリを使って

import itertools
for A in itertools.permutations(range(1, 6 + 1), 3):
    print(A)

と書くことができるが、ジェネレーターを使って、

def generate(itr, r):
    if r == 0:
        yield []
    else:
        for A in generate(itr, r - 1):
            for i in itr:
                if i not in A:
                    yield A + [i]

for A in generate(range(1, 6 + 1), 3):
    print(A)

と書くこともできる。これは以下のように考えればよい。
$1$ から $n$ までの自然数から $ r $ 個を選んで並べてできる数列というのは、$1$ から $n$ までの自然数からなる長さ $r$ の数列で、重複がないものである。
ここで、$1$ から $n$ までの自然数からなる長さ $r - 1$ の数列がすべて列挙できたとすると、それぞれの数列 $A$ の末尾に $A$ に現れない要素を追加した数列を列挙すればよい。

8クイーン問題

チェス盤に $8$ 個のクイーンを、どのクイーンも他のクイーンに取られないように置く配置を全て求めたい。
この問題を言い換えると、$0, \ldots, 7$ からなる長さ $8$ の数列 $A$ で、$i \neq j$ ならば $A _ i \neq A _ j$ かつ $| A _ i - A _ j | \neq | i - j |$ をみたすものを求めるということになり、以下のように書ける。

def generate(n=8, m=8):
    if n == 0:
        yield []
    else:
        for A in generate_queen(n - 1, m):
            for b in range(m):
                flag = True
                for i, a in enumerate(A):
                    if a == b or abs(a - b) == abs(i - (n - 1)):
                        flag = False
                if flag:
                    yield A + [b]

$0, \ldots, 7$ からなる長さ $n - 1$ で条件をみたす数列 $A$ に対して、その末尾に $b$ を追加したいとき、各 $a \in A$ とその位置 $i$ について $a=b$ または $|a-b| = | i - (n - 1) |$ となっていれば追加できず、そうでない場合に追加することで長さ $n - 1$ で条件をみたす配列ができる。

Panasonic 2020 D - String Equivalence

長さ $n$ の標準形は、長さ $n - 1$ の標準形 $S$ に、$S$ に含まれる文字か $S$ に含まれない文字のうち辞書順で最初のものを付けたものなので、数列に言い換えると、長さ $n$ の数列は、長さ $n - 1$ の数列 $A$ に $0$ から $\mathrm{max}\ A + 1$ までの数字を追加してできるものなので、以下のように書ける。

def generate(n):
    if n == 1:
        yield [0]
    else:
        for A in generate(n - 1):
            for i in range(max(A) + 2):
                yield A + [i]

ABC 161 D - Lunlun Number

先頭は $0$ でなく、$0$ から $9$ の数字からなる長さ $n$ で隣り合う数字の差が $1$ 以下の数列を列挙したい。条件をみたす長さ $n - 1$ の数列 $A$ に対して、$A$ の末尾を $x$ とすると、$|x - i| \leq 1$ なる $i$ を $A$ に追加してできる数列が条件をみたす長さ $n$ の数列である。

def generate(n):
    if n == 1:
        for i in range(1, 10):
            yield [i]
    else:
        for A in generate(n - 1):
            for i in range(10):
                if abs(A[-1] - i) <= 1:
                    yield A + [i]