Skip to main content

netoptim_rs/
dijkstra.rs

1//! Dijkstra's shortest path algorithm implementation.
2
3use petgraph::algo::FloatMeasure;
4#[allow(unused_imports)]
5use petgraph::graph::NodeIndex;
6use petgraph::visit::{
7    EdgeRef, IntoEdges, IntoNodeIdentifiers, NodeCount, NodeIndexable, VisitMap, Visitable,
8};
9use std::cmp::Ordering;
10use std::collections::BinaryHeap;
11
12/// Result of Dijkstra's shortest path algorithm.
13///
14/// Contains the distances from the source node to all other nodes,
15/// and the predecessor of each node along the shortest path.
16#[derive(Debug, Clone)]
17pub struct DijkstraResult<NodeId, EdgeWeight> {
18    pub distances: Vec<EdgeWeight>,
19    pub predecessors: Vec<Option<NodeId>>,
20}
21
22/// State for the priority queue in Dijkstra's algorithm.
23/// Contains a node and its current cost from the source.
24#[derive(Clone, Debug)]
25struct State<NodeId, Cost> {
26    node: NodeId,
27    cost: Cost,
28}
29
30impl<NodeId: PartialEq, Cost: PartialEq> PartialEq for State<NodeId, Cost> {
31    fn eq(&self, other: &Self) -> bool {
32        self.node == other.node && self.cost == other.cost
33    }
34}
35
36impl<NodeId: PartialEq, Cost: PartialEq> Eq for State<NodeId, Cost> {}
37
38impl<NodeId: PartialEq, Cost: FloatMeasure> Ord for State<NodeId, Cost> {
39    fn cmp(&self, other: &Self) -> Ordering {
40        // Use partial_cmp and unwrap since FloatMeasure guarantees total ordering via NaN handling
41        other
42            .cost
43            .partial_cmp(&self.cost)
44            .unwrap_or(Ordering::Equal)
45    }
46}
47
48impl<NodeId: PartialEq, Cost: FloatMeasure> PartialOrd for State<NodeId, Cost> {
49    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
50        Some(self.cmp(other))
51    }
52}
53
54/// \[Generic\] Compute shortest paths from node `source` to all other nodes using Dijkstra's algorithm.
55///
56/// This implementation uses a binary heap for efficient priority queue operations.
57/// The algorithm requires non-negative edge weights.
58///
59/// # Arguments
60/// * `g` - The graph to compute shortest paths on
61/// * `source` - The source node index
62///
63/// # Returns
64/// * `Ok(DijkstraResult)` - Contains distances and predecessors for each node
65/// * `Err` - If a negative edge weight is found
66///
67/// # Example
68/// ```rust
69/// use petgraph::Graph;
70/// use petgraph::prelude::*;
71/// use netoptim_rs::dijkstra::dijkstra;
72///
73/// let mut g = Graph::new();
74/// let a = g.add_node(()); // node with no weight
75/// let b = g.add_node(());
76/// let c = g.add_node(());
77/// let d = g.add_node(());
78/// g.extend_with_edges(&[
79///     (0, 1, 2.0),
80///     (0, 3, 4.0),
81///     (1, 2, 1.0),
82///     (2, 4, 5.0),
83///     (3, 4, 1.0),
84/// ]);
85///
86/// let result = dijkstra(&g, a);
87/// assert!(result.is_ok());
88/// let paths = result.unwrap();
89/// assert_eq!(paths.distances[a.index()], 0.0);
90/// assert_eq!(paths.distances[b.index()], 2.0);
91/// assert_eq!(paths.distances[c.index()], 3.0);
92/// ```
93pub fn dijkstra<G>(
94    g: G,
95    source: G::NodeId,
96) -> Result<DijkstraResult<G::NodeId, G::EdgeWeight>, String>
97where
98    G: NodeCount + IntoNodeIdentifiers + IntoEdges + NodeIndexable + Visitable,
99    G::EdgeWeight: FloatMeasure,
100{
101    let ix = |i| g.to_index(i);
102    let node_count = g.node_count();
103
104    let mut distances = vec![<_>::infinite(); node_count];
105    let mut predecessors = vec![None; node_count];
106    let mut visited = g.visit_map();
107
108    distances[ix(source)] = <_>::zero();
109
110    let mut heap = BinaryHeap::new();
111    heap.push(State {
112        node: source,
113        cost: <_>::zero(),
114    });
115
116    while let Some(State { node, cost }) = heap.pop() {
117        if visited.is_visited(&node) {
118            continue;
119        }
120
121        visited.visit(node);
122
123        for edge in g.edges(node) {
124            let target = edge.target();
125            let weight = *edge.weight();
126
127            if weight < <_>::zero() {
128                return Err("Dijkstra's algorithm requires non-negative edge weights".to_string());
129            }
130
131            let new_cost = cost + weight;
132            if new_cost < distances[ix(target)] {
133                distances[ix(target)] = new_cost;
134                predecessors[ix(target)] = Some(node);
135                heap.push(State {
136                    node: target,
137                    cost: new_cost,
138                });
139            }
140        }
141    }
142
143    Ok(DijkstraResult {
144        distances,
145        predecessors,
146    })
147}
148
149/// \[Generic\] Compute shortest path from `source` to `target` using Dijkstra's algorithm.
150///
151/// # Arguments
152/// * `g` - The graph to compute shortest path on
153/// * `source` - The source node index
154/// * `target` - The target node index
155///
156/// # Returns
157/// * `Some(path)` - A vector of node indices representing the shortest path from source to target
158/// * `None` - If no path exists
159///
160/// # Example
161/// ```rust
162/// use petgraph::Graph;
163/// use petgraph::prelude::*;
164/// use netoptim_rs::dijkstra::dijkstra_path;
165///
166/// let mut g = Graph::new();
167/// let a = g.add_node(());
168/// let b = g.add_node(());
169/// let c = g.add_node(());
170/// g.extend_with_edges(&[(0, 1, 2.0), (1, 2, 3.0)]);
171///
172/// let path = dijkstra_path(&g, a, c);
173/// assert_eq!(path, Some(vec![a, b, c]));
174/// ```
175pub fn dijkstra_path<G>(g: G, source: G::NodeId, target: G::NodeId) -> Option<Vec<G::NodeId>>
176where
177    G: NodeCount + IntoNodeIdentifiers + IntoEdges + NodeIndexable + Visitable,
178    G::EdgeWeight: FloatMeasure,
179    G::NodeId: PartialEq,
180{
181    let ix = |i| g.to_index(i);
182
183    let result = dijkstra(g, source).ok()?;
184
185    if result.predecessors[ix(target)].is_none() && source != target {
186        return None;
187    }
188
189    let mut path = vec![target];
190    let mut current = target;
191
192    while current != source {
193        let pred = result.predecessors[ix(current)]?;
194        path.push(pred);
195        current = pred;
196    }
197
198    path.reverse();
199    Some(path)
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205    use petgraph::Graph;
206
207    #[test]
208    fn test_dijkstra_simple() {
209        let mut g = Graph::new();
210        let a = g.add_node(());
211        let _b = g.add_node(());
212        let _c = g.add_node(());
213        let _d = g.add_node(());
214        g.extend_with_edges([(0, 1, 2.0), (0, 3, 4.0), (1, 2, 1.0), (3, 4, 1.0)]);
215
216        let result = dijkstra(&g, a);
217        assert!(result.is_ok());
218        let paths = result.unwrap();
219        assert_eq!(paths.distances[0], 0.0);
220        assert_eq!(paths.distances[1], 2.0);
221        assert_eq!(paths.distances[2], 3.0);
222        assert_eq!(paths.distances[3], 4.0);
223    }
224
225    #[test]
226    fn test_dijkstra_negative_weight() {
227        let g: Graph<(), f32> = Graph::<(), f32>::from_edges([(0, 1, -1.0)]);
228        let result = dijkstra(&g, NodeIndex::new(0));
229        assert!(result.is_err());
230    }
231
232    #[test]
233    fn test_dijkstra_path() {
234        let mut g = Graph::new();
235        let a = g.add_node(());
236        let b = g.add_node(());
237        let c = g.add_node(());
238        g.extend_with_edges([(0, 1, 2.0), (1, 2, 3.0)]);
239
240        let path = dijkstra_path(&g, a, c);
241        assert_eq!(path, Some(vec![a, b, c]));
242    }
243
244    #[test]
245    fn test_dijkstra_path_no_path() {
246        let mut g = Graph::new();
247        let a = g.add_node(());
248        let b = g.add_node(());
249        let c = g.add_node(());
250        g.add_edge(a, b, 1.0);
251
252        let path = dijkstra_path(&g, a, c);
253        assert_eq!(path, None);
254    }
255
256    #[test]
257    fn test_dijkstra_disconnected() {
258        let mut g: Graph<(), f64> = Graph::new();
259        let a = g.add_node(());
260        let b = g.add_node(());
261        let c = g.add_node(());
262        g.add_edge(a, b, 1.0);
263
264        let result = dijkstra(&g, a);
265        assert!(result.is_ok());
266        let paths = result.unwrap();
267        assert_eq!(paths.distances[a.index()], 0.0);
268        assert_eq!(paths.distances[b.index()], 1.0);
269        assert!(paths.distances[c.index()].is_infinite());
270    }
271
272    #[test]
273    fn test_dijkstra_same_node() {
274        let mut g: Graph<(), f64> = Graph::new();
275        let a = g.add_node(());
276
277        let result = dijkstra(&g, a);
278        assert!(result.is_ok());
279        let paths = result.unwrap();
280        assert_eq!(paths.distances[a.index()], 0.0);
281    }
282
283    #[test]
284    fn test_dijkstra_path_same_node() {
285        let mut g: Graph<(), f64> = Graph::new();
286        let a = g.add_node(());
287
288        let path = dijkstra_path(&g, a, a);
289        assert_eq!(path, Some(vec![a]));
290    }
291}