ぼくだけのDijkstraライブラリを紹介するよ

この記事は ProgatePathコミュニティ Qiitaアドベントカレンダー 10日目の記事です。

これはなに？

前提として、ダイクストラ法はご存知でしょうか。ダイクストラ法について知らないよ！という人向けに簡単に説明をすると、重み付きのグラフに対して、ある単一始点から各ノードへ移動するための最短距離を高速に求めるアルゴリズムです。詳しくは ダイクストラ法による単一始点最短経路を求めるアルゴリズム - アルゴリズムロジック などが参考になると思います。

計算量

よくあるダイクストラ法の実装では計算量は頂点数 \(V\) , 辺の本数 \(E\) に対して \(O(E \log{V})\) です。これは、候補となる頂点の取得・更新はHeapを用いることで \(O(\log{N})\) で行え、辺について候補の作成のためにすべてを見る必要があるからです。

本題

ダイクストラ法はグラフによっては大きく定数倍高速化できる場合があります。グラフの条件はかなり簡単に言うと「始点から他頂点までの最短経路のうち、距離が最も長いものよりも重みが大きい辺が多数存在する」という場合です。
以下は、ナイーブなダイクストラ法の実装例です。startを受け取ってその頂点から他の頂点への最短距離を返却します。

fn dijkstra(g: &Vec<Vec<(usize, u64)>>, start: usize) -> Vec<u64> {
    let n = g.len();
    let mut dist = vec![INF; n];

    let mut heap = BinaryHeap::new();
    heap.push(Reverse((0, start)));

    while let Some(Reverse((cost, pos))) = heap.pop() { 
        if dist[pos] != INF { 
            continue;
        }
        dist[pos] = cost;
        for &(ni, w) in &g[pos] {
            if dist[ni] == INF {
                heap.push(Reverse((cost + w, ni)));
            }
        }
    }
    dist
}

このようなナイーブなコードは先に上げたような条件下では無駄が多くなります。これは、すべての頂点に対する最短経路が求まったにもかかわらず、heapに多数の辺が残っているためにその辺をすべて取り出すまでループを繰り返すためです。
ということで、このボトルネックを解消すればよいです。ループごとに「全頂点に到達したか？」を素直に判定すると追加で \(O(V)\) かかってしまうので、代わりに「まだ到達していない頂点が幾つあるか？」を保持して、これが0になったらbreakすることにします。こうすれば \(O(1)\) でできて嬉しいですね。改善版のコードを以下に示します。

fn dijkstra(g: &Vec<Vec<(usize, u64)>>, start: usize) -> Vec<u64> {
    let n = g.len();
    let mut dist = vec![INF; n];
    let mut left = n;
    let mut heap = BinaryHeap::new();
    heap.push(Reverse((0, start)));

    while let Some(Reverse((cost, pos))) = heap.pop() {
        if dist[pos] != INF {
            continue;
        }
        dist[pos] = cost;
        left -= 1;
        if left == 0 {
            break;
        }

        for &(ni, w) in &g[pos] {
            if dist[ni] == INF {
                heap.push(Reverse((cost + w, ni)));
            }
        }
    }
    dist
}

実際どれぐらい改善するの？

ABC369E - Sightseeing Tour を例に、以下の2つの提出を見比べていただくのが早いと思います。

AC(1939ms) / TLE


そういえばログインしないと提出詳細を見れなかったような気がしたので、例としてテストケースから1つ結果を持ってきました。
画像上が定数倍高速化を用いたAC解で、画像下が定数倍高速化を使わないTLE解です。このケースだけ見ても、約3倍もの高速化に成功しています！

おしまい

いかがでしたか？(定型) 意外とこういう事している人は少ない ので、ぼくのライブラリの微妙な工夫を紹介してみました。誰得なんだ、これ
一応釈明しておくと普段はちゃんと使いやすいようにライブラリ整備をしています。かなり使いやすく整備しているので、気になったら見ていってください。
GitHub - dijkstra.rs(ライブラリ本体)
GitHub - shortest_path.rs(使用例)

なお、普通のグラフに対しては分岐 が入るので通常のものよりも遅くなる場合もあります(え？)。とはいえ、影響は普通に使う分には数ms程度に収まる上に、キラーケースに対して強くなってほしいな〜という気持ちを込めてぼくは採用しています。