幅優先探索の亜種ですね。
アルゴリズムの流れは下記です。
- rootノードのみからなるpath、を持つpathsリストを生成(paths = [ [root] ])
- paths内の全てのpathを1階層分だけ展開しpathsを更新
- スコアが最も高いk個のpathのみをpathsに残す
- 2に戻る
補足: https://www.baeldung.com/cs/beam-search
下記のように実装しました。
コード
Notesに書いてあるような次の関数を持つapiを作って実行すると動きます。
- init : 初期化関数 api.init() がbeam_searchの初めに1度呼ばれます
- step : pathを入力として探索を1階層進めた pathのiterator/generatorを返す関数です
- score : pathを入力としてそのpathの評価値を返す関数です。値が高いほど良いことを表します。
- count : 探索が1ラウンド(上記の説明の3)が終了すると呼び出されます。
- terminate : beam_searchを終了するべきであればtrue, そうでないならfalseを返します。
ヒープの管理はpython標準のheapqを用いています。また今回はheapの大きさが固定長なので、heapの大きさがkになったらheapq.pushpop関数を用いてpushとpopを一つの関数で実行しています。(別々にやるよりもこの関数を使う方が効率が良いらしい; 公式ドキュメントより)
例
TSP(traveling salesman problem)を例に実行してみます。
class TSPAPI: def __init__(self, n): self.iter = 0 self.n = n def init(self): self.iter = 0 def step(self, path): for i in range(self.n): if i not in path: yield path[:] + [i] def score(self, path): """total distance of path""" return -sum(D[i][j] for i, j in zip(path, path[1:])) def count(self): self.iter += 1 def terminate(self): return self.iter >= self.n # TSP N = 4 D = [ [0, 1, 2, 3], [4, 0, 2, 1], [1, 2, 0, 3], [2, 3, 4, 0], ]
Nが都市数。Dが移動コスト行列です。
- step(path)ではまだ訪問していない都市(i)をpathの最後に加えて返します
- score(path)ではpathの巡回コストを出力
- count()では訪問都市数のカウントをして
- terminate()では訪問都市数が全都市数になったタイミングでtrueを出すようにしています
さて、このコードを k=1で実行してみます。
from beam_search import beam_search root = [] k = 1 api = TSPAPI(N) paths, scores = beam_search(root, k, api) for path, score in zip(paths[::-1], scores[::-1]): print("path", path, "distance", -score)
path [3, 0, 1, 2] distance 5
k=2で実行してみます。
from beam_search import beam_search root = [] k = 3 api = TSPAPI(N) paths, scores = beam_search(root, k, api) for path, score in zip(paths[::-1], scores[::-1]): print("path", path, "distance", -score)
path [2, 0, 1, 3] distance 3 path [3, 0, 1, 2] distance 5
ということでk = 2の方がより少ないコストのパス([2, 0, 1, 3])をつけることができていますね。
まとめ
beam searchの実装を行いました。