抽象化されたBFSと復元

最近はライブラリ盆栽が楽しいのでいろいろ作ってるよのいろいろの部分です。
すべてを書き終えてから実装の話しかしていないことに気が付きましたが、抽象化の真面目な話は@naoya_itoさんのスライド「Haskell でアルゴリズムを抽象化する / 関数型言語で競技プログラミング」が大変参考になるかと思います。
競プロでBFSを書くとき、大体以下のような実装をすると思います。
while let Some(u) = que.pop_front() {
for v in &g[u] {
if seen[v] {
continue;
}
seen[v] = true;
...
}
}
これ自体はまぁそう、という感じではあるんですが、たとえばグリッドだったら境界チェックが入ったり、問題固有の制約があったりというのを考えて実装に落としていく仮定で、このような「いつもやる処理」隠蔽して、問題固有の部分だけ記述すれば勝手にBFSして答えが帰ってくるねという状況にしたいという気持ちになります。復元パートとかは書くのが面倒だったりしますし。
で、作ったはいいんですが使い方を忘れていることが多いのでメモも兼ねての記事です。
使い方
以下のトレイトを実装すればよいです。
pub trait BfsHandler {
type State;
fn neighbors(&mut self, state: &Self::State) -> Vec<Self::State>;
fn mark_visited(&mut self, state: &Self::State, from: Option<&Self::State>);
fn is_visited(&self, state: &Self::State) -> bool;
fn on_visited(&mut self, _state: &Self::State) -> bool {
true
}
fn should_stop(&self) -> bool {
false
}
fn parent(&self, _state: &Self::State) -> Option<Self::State> {
None
}
}
State
「今の状態」を表す型です。たとえばグラフの探索だったらいまの頂点という気持ちでusize を埋めるでしょうし、ゲームの状態探索だったらVec<T> だったりするかもしれません。
neighbors
「今の状態から一手で到達可能な状態の集合」を返します。集合と言っていますがまぁ別に多重集合になっていてもよいです(それは普通にBFS書いているときに気にしないのと同じ)。
is_visited
名前の通り、その状態に訪れたことがあるかどうかを返します。trueが返された場合にはその頂点からのneighborsの列挙を行いません。conitnue;するのと同じと思えば良いでしょう。
on_visited
その状態に到達したときに行う処理で、主に枝刈りを想定して用意しています。falseを返すとconitnue;するのと同じです。が、一度もつかった記憶がありません。デフォルトでtrueを返すので枝刈り不要なら書かなくてよいです。
shoud_stop
trueを返すと探索の強制打ち切りが出来ます。いらないかもしれないが、定数倍が怖いときのお守り程度にはなるかもしれません。
parent
名の通り、状態が与えられてその親の状態が何だったかを返します。復元用途で、不要ならNoneを返せばよいです。
これらの実装をした後、bfs(handler, starts); を呼ぶだけで簡単にBFSができます。startsは始点の集合で、多始点BFS用途。復元が必要であればbfs後にrestore_path(handler, end);を呼ぶだけです。
実装
実装はここ。まあ特段変なことをやっているわけではなく、いつものBFSの中で上記の実装を適宜呼び出しているだけです。復元も始点がNoneであるという前提を持たせて、Noneが出てくるまでParentを掘るという至って自然な?方法を取っています。AHCとかで無限回やるやつ。
実際につかった例として、ABC453 - Dの提出を乗せておきます。復元つき。提出(Rust, 912ms)
pub fn bfs<H: BfsHandler>(handler: &mut H, starts: impl IntoIterator<Item = H::State>) {
let mut que = VecDeque::new();
for start in starts {
if handler.is_visited(&start) {
continue;
}
handler.mark_visited(&start, None);
que.push_front(start);
}
while let Some(state) = que.pop_front() {
if handler.should_stop() {
break;
}
if !handler.on_visited(&state) {
continue;
}
for nxt in handler.neighbors(&state) {
if handler.is_visited(&nxt) {
continue;
}
handler.mark_visited(&nxt, Some(&state));
que.push_back(nxt);
}
}
}
pub fn restore_path<H>(handler: &H, target: H::State) -> Vec<H::State>
where
H: BfsHandler,
H::State: Clone,
{
let mut path = vec![target.clone()];
let mut cur = target;
while let Some(p) = handler.parent(&cur) {
path.push(p.clone());
cur = p;
}
path.reverse();
path
}
注意点として、neighbors がVec<Self::State> を返す都合で必ずheap allocationが走るという問題があります。AtCoder環境ではPythonで通るぐらいの制約にはなっているし、最近ぼくが海外コンに全然出ていないので大きく問題になることはまぁないことから放置していますが、まあまあ遅いです。さっきの問題で900msぐらい、持ってる配列をフラットにしても600~700msぐらいだった気がする。解決策自体はあるものの、設計としてあまり嬉しくないので放置しています。この辺は良い解決策を模索したい気持ち…。
いつものパートを書かなくて良くなるという抽象化の恩恵が果たして今の自分にどれぐらい効くかはよくわからないですが、割と書いていて楽しいので(?)おすすめです。