2日目: Dijkstra法を実装する

競プロ用ライブラリを作る Advent Calendar 2024の2日目です.

Dijkstra法って?

辺のコストが非負の有向グラフが与えられたとき, ある頂点からの最短経路を$O((E+V)\log V)$などのオーダーで計算するアルゴリズムです.

仕組み

ざっくり「近い順に距離を確定していく」という方法です. とりあえず距離が確定したかどうかの情報が入る配列knownを作りましょう

fn sample_dijkstra(edges: &[Vec<(usize, u32)>], start: usize, goal: usize) -> Option<u32> {
    let n = edges.len();
    // 最初はどの頂点も距離が決定していないのでfalseを詰める
    let mut known = vec![false; n];

    // 略
}

そして, スタート地点は距離$0$で確定してます. 自明ですね.

fn sample_dijkstra(edges: &[Vec<(usize, u32)>], start: usize, goal: usize) -> Option<u32> {
    // 略

    known[start] = true;

    // 略
}

さて, ここで距離が確定していない頂点の中で「距離が確定した頂点から直接行ける最も近い頂点」を探すと, その頂点の最短距離は直接行った時の距離になります.

なぜかというと...

この方法で選んだ頂点がより短い経路を持っていたと仮定し, 矛盾を導きます.

スタート地点は距離が確定していてゴール地点は距離が確定していないので, どこかに「距離が確定している頂点から距離が確定していない頂点への辺」があります

スタートからその辺までのパスを抜き出してやると, これは「距離が確定した頂点から距離が確定していない頂点へ直接行くパス」で, 最初に選んだものより短いと仮定した経路の一部分なので, 最初に選んだものより短いです.
これは最初に条件を満たす最も近い頂点を選んだことと矛盾します.

背理法から仮定は正しくなく, この方法でとった頂点の距離は最短経路になっていることが分かりました.

そんな感じで目的の頂点の距離が分かったり, 到達できる頂点の距離が全てわかったりするまで「距離が確定した頂点から直接行ける最も近い頂点」を追加し続けることで正しい距離が計算できます.

ここで,「距離が確定した頂点から直接行ったときの距離」は毎回計算するのではなく計算結果を使いまわすと高速化出来ます. ある頂点の距離を確定したとき, その距離を確定した頂点から直接行ける頂点での値だけが変わるので, そういう計算結果を保持するcost_upperを作ります.

fn sample_dijkstra(edges: &[Vec<(usize, u32)>], start: usize, goal: usize) -> Option<u32> {
    // 略

    let mut costs_upper = vec![None; n];

    // スタート地点から直接行ける頂点の情報を反映する
    for &(v, weight) in edges[start].iter() {
        costs_upper[v] = Some(weight);
    }

    // 略
}

そしたら一番近い頂点を確定する作業を繰り返します.

fn sample_dijkstra(edges: &[Vec<(usize, u32)>], start: usize, goal: usize) -> Option<u32> {
    // 略

    loop {
        // 距離が確定してない中で一番近い頂点を v, そのコストを cost とする
        let Some((cost, v)) = costs_upper
            .iter()
            .enumerate()
            // 距離が確定していない頂点だけ抜き出す
            .filter(|&(index, _)| !known[index])
            // 距離が確定している頂点から直接行けてた頂点だけ抜き出す
            .filter_map(|(index, &cost)| cost.map(|cost| (cost, index)))
            // 距離が一番近いものを見つける
            .min()
        else {
            // 条件を満たす頂点がもうなかったら「到達不可能」ということでNoneを返す
            return None;
        };
        // ゴールの距離が確定したらその距離を返す
        if v == goal {
            return Some(cost);
        }
        // 距離が確定した頂点の情報を更新する
        known[v] = true;
        // 頂点vから直接行ける頂点を全部見る
        for &(u, weight) in edges[v].iter() {
            let cost = cost + weight;
            if costs_upper[u].is_none_or(|c| c > cost) {
                costs_upper[u] = Some(cost);
            }
        }
    }
}

さて, この実装の計算量は$O(V^2)$です.
この計算のボトルネックは「毎回全ての頂点の中から一番近い頂点を探している」というところです.
これは優先度付きキューを用いれば解決できます. costs_upperの中身を全部std::collections::Binaryheapなどに入れてやりましょう.
値を更新したとき, このままだと更新後の要素を入れるだけでなく更新前の要素を削除する必要が出てきますが, その操作をするのは非常に面倒なので, costs_upperの値と比較して違ってたら無視するようにすれば問題ないですね.
std::collections::BinaryHeapは値を大きい方から返すので気を付けてください. 競プロではコストをビット反転して入れてやるのがお手軽です.

fn sample_dijkstra2(edges: &[Vec<(usize, u32)>], start: usize, goal: usize) -> Option<u32> {
    use std::collections::BinaryHeap;
    let n = edges.len();
    let mut known = vec![false; n];
    known[start] = true;

    let mut costs_upper = vec![None; n];
    // 2分ヒープの作成
    let mut heap = BinaryHeap::new();

    for &(v, weight) in edges[start].iter() {
        costs_upper[v] = Some(weight);
        // ヒープにも値を入れていく
        heap.push((!weight, v));
    }
    while let Some((cost, v)) = heap.pop() {
        // ビット反転した数値を入れているので元に戻してやる
        let cost = !cost;
        // コストがcost_upperに入ってる値と違うなら無視する
        if costs_upper[v] != Some(cost) {
            continue;
        }
        if v == goal {
            return Some(cost);
        }
        known[v] = true;
        for &(u, weight) in edges[v].iter() {
            let cost = cost + weight;
            if !known[u] && costs_upper[u].is_none_or(|c| c > cost) {
                costs_upper[u] = Some(cost);
                // 条件を満たすなら追加する
                heap.push((!cost, u));
            }
        }
    }
    None
}

これで計算量は$O(E\log V)$になります.

このままでも問題ないですが, もう少し簡潔な書き方があるのでします.
実は配列knownは最早必要ありません. 現状この配列を参照しているのはcosts_upperを更新しても良いか判定する部分ですが, 一緒に「距離を更新可能か」の判別もしています. 距離が確定していた場合は必ずこの判定に引っかかるので, knownは使う必要がなくなり, 消せます.

fn sample_dijkstra3(edges: &[Vec<(usize, u32)>], start: usize, goal: usize) -> Option<u32> {
    use std::collections::BinaryHeap;
    let n = edges.len();
    let mut costs_upper = vec![None; n];
    let mut heap = BinaryHeap::new();

    for &(v, weight) in edges[start].iter() {
        costs_upper[v] = Some(weight);
        heap.push((!weight, v));
    }
    while let Some((cost, v)) = heap.pop() {
        let cost = !cost;
        if costs_upper[v] != Some(cost) {
            continue;
        }
        if v == goal {
            return Some(cost);
        }
        for &(u, weight) in edges[v].iter() {
            let cost = cost + weight;
            if costs_upper[u].is_none_or(|c| c > cost) {
                costs_upper[u] = Some(cost);
                heap.push((!cost, u));
            }
        }
    }
    None
}

これが競プロでよく使われる形のDijkstra法です.

実装

• dijkstra.rs

Dijkstra法では最後に見た頂点より距離が小さい頂点が入ることは無いので, 前日に実装したRadixHeapが使えます.

まとめ

いかがでしたか?
明日はみんな大好きUnionFindを実装します.