Skip to content

Bfs

Rust

mod graph {
    use std::{
        cmp::Reverse,
        collections::{BinaryHeap, VecDeque},
    };

    use num::Zero;

    #[derive(Clone, Copy)]
    struct Edge<T> {
        from: usize,
        to: usize,
        cost: T,
    }

    pub struct Graph<T: Copy> {
        n: usize,
        edges: Vec<Vec<Edge<T>>>,
    }

    impl<T: Copy> Graph<T> {
        pub fn new(n: usize) -> Self {
            let edges = vec![vec![]; n];
            Self { n, edges }
        }

        pub fn add_directed_edge(&mut self, from: usize, to: usize, cost: T) {
            self.edges[from].push(Edge { from, to, cost })
        }

        pub fn add_edge(&mut self, from: usize, to: usize, cost: T) {
            self.add_directed_edge(from, to, cost);
            self.add_directed_edge(to, from, cost);
        }

        /// 頂点 s から各頂点への最短距離を bfs で計算する
        /// 全ての辺の長さは 1 と仮定する
        pub fn bfs(&self, s: usize) -> Vec<Option<u32>> {
            let mut dist = vec![None; self.n];
            let mut q = VecDeque::with_capacity(self.n);

            dist[s] = Some(0);
            q.push_back(s);
            while let Some(v) = q.pop_front() {
                for edge in &self.edges[v] {
                    if dist[edge.to].is_some_and(|c| c <= dist[edge.from].unwrap() + 1) {
                        continue;
                    }
                    dist[edge.to] = Some(dist[edge.from].unwrap() + 1);
                    q.push_back(edge.to);
                }
            }

            dist
        }
    }

    impl<T: Ord + Copy + Zero> Graph<T> {
        /// 頂点 s から各頂点への最短距離をダイクストラ法で計算する
        pub fn dijkstra(&self, s: usize) -> Vec<Option<T>> {
            let mut dist = vec![None; self.n];
            let mut q = BinaryHeap::<Reverse<(T, usize)>>::new();

            dist[s] = Some(T::zero());
            q.push(Reverse((T::zero(), s)));
            while let Some(Reverse((cost, v))) = q.pop() {
                if dist[v].is_some_and(|c| c < cost) {
                    continue;
                }

                for edge in &self.edges[v] {
                    let new_cost = cost + edge.cost;
                    if dist[edge.to].is_some_and(|c| c <= new_cost) {
                        continue;
                    }
                    dist[edge.to] = Some(new_cost);
                    q.push(Reverse((new_cost, edge.to)));
                }
            }

            dist
        }
    }
}