1use 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#[derive(Debug, Clone)]
17pub struct DijkstraResult<NodeId, EdgeWeight> {
18 pub distances: Vec<EdgeWeight>,
19 pub predecessors: Vec<Option<NodeId>>,
20}
21
22#[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 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
54pub 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
149pub 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}