Rustのpetgraphクレートを使ってAtCoderのグラフ問題を解く

はじめに

AtCoderの言語アップデートによりRustでいろいろなクレート(ライブラリ)が使えるようになりました。

ここではグラフアルゴリズムを提供するpetgraphクレートを使ってグラフの問題を解いてみます。

ABC 070 D - Transit Tree Path

問題概要
n頂点の木とQ個のクエリと整数Kが与えられます。
j(1<=j<=Q)番目のクエリ:頂点x_jから頂点Kを経由した頂点y_jまでの最短距離を求めよ

制約概要
3<= N <= 10^5
1<= Q <= 10^5
1 <= K <= N
1 <= x_j, y_j <= N(1 <= j <= Q)
x_j != y_j(1<= j <= Q)
x_j != K, y_j != K(1 <= j <= Q)

use petgraph::algo::dijkstra;
use petgraph::graph::{NodeIndex, UnGraph};
use proconio::input;

fn main() {
    input! {
        n: usize,
        edges: [(usize, usize, usize); n-1],
        q: usize,
        k: usize,
        pairs: [(usize, usize); q],
    }

    //グラフを作る。
    let g = UnGraph::<usize, usize, usize>::from_edges(&edges);

    //ダイクストラ法で頂点kから他の頂点の最短距離を計算
    let res = dijkstra(&g, k.into(), None, |e| *e.weight());


    for (start, end) in &pairs {
        //頂点x_jと頂点kの最短距離 + 頂点kと頂点y_jの最短距離を計算
        let distance =
            res.get(&NodeIndex::new(*start)).unwrap() + res.get(&NodeIndex::new(*end)).unwrap();
        println!("{}", distance);
    }
}

解法
頂点x_jから頂点Kを経由した頂点y_jまでの最短距離というのは頂点x_jと頂点Kの最短距離+頂点kと頂点y_jの最短距離と等しいです。したがって事前に頂点Kから他の頂点の最短距離を求めておけば各クエリはO(1)で処理できます。

おわりに

petgraphクレートを使うことでグラフとアルゴリズムを一切実装せず、解くことができました。petgraphクレートにはほかにも様々なデータ構造やアルゴリズムが存在するので使いこなせれば、力強い味方になってくれそうな気がします。