RustでA*を書く - じゃあ、おうちで学べる

はじめに

AI時代でも、アルゴリズムや基礎を学んでおこうと若い人に話しています。そう勧める自分の復習も兼ねて、今回はRustでA*を実装します。「最短の経路を探して」とAIへ頼む時に、自分は何を最小にしたいのか。返ってきた答えを最小だと言える条件まで、小さな例で確かめます。

ある作業の担当者が、完了までの手順を選ぶ場面を考えます。2つの操作で済む案と、3つの操作を順に進める案があります。操作数で選ぶなら前者です。ところが所要時間を足すと、前者は9分、後者は3分です。待っている人へ早く結果を返すには、操作の回数より、各操作にかかる時間を比べる必要があります。

候補が増えると、どの手順から調べるかも問題になります。そこまでにかかる時間に、残り時間の見積もりを加えれば、早く終わりそうな候補を先に調べられます。ただし、見積もりのせいで本当は早い候補を後回しにし、探索を終えてしまっては困ります。AIへ手順の選択を頼む時も、何を最小にしたいかと、見積もりに求める条件が関わります。

A*は、移動費用の合計が最小になる経路を探す方法です。各地点への最良コストを記録しながら、残りの費用も見積もって探索順を決めます。この記事ではさらに、未到達を整数の最大値で表すと、同じ費用で到達する経路を区別できなくなる問題を扱います。探索順と記録の表し方を、両方確かめます。

ここまでの費用と残りの見積もりを足す

移動回数が少なくても、費用の高い辺を通れば合計は高くなります。A*では、出発点からの実費と、ゴールまでに必要な残り費用の見積もりを足して、次に調べる候補を選びます。

すでに通った辺の費用の合計をg、残りの見積もりをhとすると、調べる優先度はf = g + hです。hをヒューリスティックと呼びます。最小コストを保証するには、実際の残りコストを超えない見積もりが必要です。a_starは非負の整数を受け取り、ゴールの見積もりは0とします。すべて0なら、そこまでの実費が小さい順で調べます。

edges[i]には「移動先の番号と辺の費用」の組を並べます。例の0→2→1→3は1 + 1 + 1 = 3で、0→1→3の9より安くなります。別の経路から安く到達できたらbestを更新し、そのノードを再び候補へ追加します。この実装の戻り値は、求めた最小費用です。

未到達を整数の最大値にしない

整数の最大値を「未到達」に使うと、その費用で着く経路を表せなくなります。 そこで最良コストをOption<u32>にし、NoneとSome(u32::MAX)を区別します。優先度はu64で計算し、経路コストがu32を超える候補は除外します。

BinaryHeapは通常、大きい値から取り出します。値をReverseで包むと、小さい優先度から取り出せます。候補には優先度、ここまでの実費、ノード番号を持たせます。checked_addは足し算がu32の範囲を超える場合にNoneを返します。

Noneを使うと、未到達の状態と、到達済みの数値を分けられます。経路コストが表現範囲を超えた時も、折り返して小さい数になることはありません。辺の費用は非負なので、上限を超えた後に移動を続けて、再び上限以内へ戻る経路もありません。

use std::{cmp::Reverse, collections::BinaryHeap};

fn a_star(edges: &[Vec<(usize, u32)>], h: &[u32], start: usize, goal: usize)
    -> Option<u32>
{
    let n = edges.len();
    if start >= n || goal >= n || h.len() != n
        || edges.iter().flatten().any(|&(v, _)| v >= n)
    { return None; }
    let mut best = vec![None; n];
    let mut queue = BinaryHeap::new();
    best[start] = Some(0);
    queue.push(Reverse((u64::from(h[start]), 0_u32, start)));
    while let Some(Reverse((_, cost, node))) = queue.pop() {
        if best[node] != Some(cost) { continue; }
        if node == goal { return Some(cost); }
        for &(next, weight) in &edges[node] {
            let Some(candidate) = cost.checked_add(weight) else { continue; };
            if best[next].is_none_or(|known| candidate < known) {
                best[next] = Some(candidate);
                queue.push(Reverse((
                    u64::from(candidate) + u64::from(h[next]), candidate, next,
                )));
            }
        }
    }
    None
}

fn main() {
    let edges = vec![vec![(1, 8), (2, 1)], vec![(3, 1)],
                     vec![(1, 1), (3, 7)], vec![]];
    assert_eq!(a_star(&edges, &[2, 1, 1, 0], 0, 3), Some(3));
    let maximum = vec![vec![(1, u32::MAX)], vec![]];
    assert_eq!(a_star(&maximum, &[0, 0], 0, 1), Some(u32::MAX));
}

4地点の例で返るのはSome(3)です。1本の辺の費用がu32::MAXなら、結果はSome(u32::MAX)になります。整数の最大値も有効な費用です。Noneは入力が不正な場合と、表せる範囲の経路がない場合に返ります。

同じノードの古い候補はキューに残るため、取り出したコストと現在の最良値を比較して捨てています。推定値が過大でないかはこの関数では検証できません。見積もりの条件を破ると、最短でないゴールを先に取り出す可能性があります。

見積もりが正解を変えてしまう例

出発点SからゴールGへ直接進む費用を3、S→AとA→Gをそれぞれ1とします。本当の最小コストは2です。ところがAからの残り費用を3と見積もれば、Aの優先度は1 + 3 = 4になります。直接Gへ進む候補の優先度3より大きいため、この実装はコスト3で探索を終えます。推定値の条件は、速さだけでなく正解にも関わります。

見積もりを使う理由と、使わない場合を比べる

作業手順を自動で選ぶシステムなら、状態をノード、作業を辺、所要時間などを費用として表せます。先ほどの例が示すように、残り時間の予測をA*へ入れる際には、過大評価が最小コストの保証を壊さないか確認が必要です。予測を0にした探索を、比較の基準として使えます。

残り費用を0とする案と、予測器の値を使う案を比べてみます。0なら見積もりの計算は不要で、非負の費用に対して過大評価もしません。ただし、残り費用による探索の絞り込みはできません。予測器を使う案では調べる候補を減らせる可能性がある一方、予測の費用と過大評価を考える必要があります。「予測がよく当たる」ことだけでは、最小費用を保証する根拠になりません。

最小費用が必須で、使える下限値もまだ分からないなら、まず0を使って探索量を測れます。処理時間が用途の上限を超えた時に、保証できる見積もりの導入を検討します。最小性を緩めるなら、その変更は利用者が必要とする答えにも関わります。探索順を改善する工夫と、答えの条件を変える判断を区別して残すことで、速くなった理由を後から説明できます。

結果を受け取った側が、次の動作を決められるか

この関数は、不正な入力と、表現できる費用の範囲で経路がない状態をNoneへまとめます。後者には、到達不能と、到達する全経路で整数の上限を超えるケースが含まれます。利用者へ理由を伝える場面では、これだけでは情報が足りません。入力を直すべき時に「経路なし」と表示しても、問題は解消しません。用途を広げるなら、入力の誤りはエラーとして分け、探索結果とも区別します。

探索に時間制限を加える場合にも、制限内で見つからなかったことを到達不能と同じに扱わない設計が必要です。途中の経路を返すなら、最小費用だと確認できたかも示します。戻り値には、答えだけでなく、その答えについて確認できた範囲を表します。

AIへ「経路がなければNoneを返して」と依頼するときも、費用が上限を超える経路をどう扱うかは残ります。経路そのものがない場合と、今回の数値型では費用を表せない場合では、呼び出す側の対応が違います。有効な費用の範囲と、失敗後に何をしたいかまで示せば、戻り値に必要な区別を検討できます。最小費用が必須なら、見つかった費用を返すだけでは完了にできないことも共有します。

検査も、費用の値と失敗時の意味を分けて考えます。ここでの実装なら、整数の最大値、到達不能、遠回りの方が安い経路を、小さな入力で既知の答えと比較できます。戻り値を細分化した後は、呼び出し側が理由に応じた動作を選べるかも確認します。