最短路径-Dijkstra

  1. use std::cmp::Reverse;
  2. use std::collections::{BTreeMap, BinaryHeap};
  3. use std::ops::Add;
  5. type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;
  7. // performs Dijsktra's algorithm on the given graph from the given start
  8. // the graph is a positively-weighted undirected graph
  9. //
  10. // returns a map that for each reachable vertex associates the distance and the predecessor
  11. // since the start has no predecessor but is reachable, map[start] will be None
  12. pub fn dijkstra<V: Ord + Copy, E: Ord + Copy + Add<Output = E>>(
  13.     graph: &Graph<V, E>,
  14.     start: &V,
  15. ) -> BTreeMap<V, Option<(V, E)>> {
  16.     let mut ans = BTreeMap::new();
  17.     let mut prio = BinaryHeap::new();
  19.     // start is the special case that doesn't have a predecessor
  20.     ans.insert(*start, None);
  22.     for (new, weight) in &graph[start] {
  23.         ans.insert(*new, Some((*start, *weight)));
  24.         prio.push(Reverse((*weight, new, start)));
  25.     }
  27.     while let Some(Reverse((dist_new, new, prev))) = prio.pop() {
  28.         match ans[new] {
  29.             // what we popped is what is in ans, we'll compute it
  30.             Some((p, d)) if p == *prev && d == dist_new => {}
  31.             // otherwise it's not interesting
  32.             _ => continue,
  33.         }
  35.         for (next, weight) in &graph[new] {
  36.             match ans.get(next) {
  37.                 // if ans[next] is a lower dist than the alternative one, we do nothing
  38.                 Some(Some((_, dist_next))) if dist_new + *weight >= *dist_next => {}
  39.                 // if ans[next] is None then next is start and so the distance won't be changed, it won't be added again in prio
  40.                 Some(None) => {}
  41.                 // the new path is shorter, either new was not in ans or it was farther
  42.                 _ => {
  43.                     ans.insert(*next, Some((*new, *weight + dist_new)));
  44.                     prio.push(Reverse((*weight + dist_new, next, new)));
  45.                 }
  46.             }
  47.         }
  48.     }
  50.     ans
  51. }
  53. #[cfg(test)]
  54. mod tests {
  55.     use super::{dijkstra, Graph};
  56.     use std::collections::BTreeMap;
  58.     fn add_edge<V: Ord + Copy, E: Ord>(graph: &mut Graph<V, E>, v1: V, v2: V, c: E) {
  59.         graph.entry(v1).or_insert_with(BTreeMap::new).insert(v2, c);
  60.         graph.entry(v2).or_insert_with(BTreeMap::new);
  61.     }
  63.     #[test]
  64.     fn single_vertex() {
  65.         let mut graph: Graph<usize, usize> = BTreeMap::new();
  66.         graph.insert(0, BTreeMap::new());
  68.         let mut dists = BTreeMap::new();
  69.         dists.insert(0, None);
  71.         assert_eq!(dijkstra(&graph, &0), dists);
  72.     }
  74.     #[test]
  75.     fn single_edge() {
  76.         let mut graph = BTreeMap::new();
  77.         add_edge(&mut graph, 0, 1, 2);
  79.         let mut dists_0 = BTreeMap::new();
  80.         dists_0.insert(0, None);
  81.         dists_0.insert(1, Some((0, 2)));
  83.         assert_eq!(dijkstra(&graph, &0), dists_0);
  85.         let mut dists_1 = BTreeMap::new();
  86.         dists_1.insert(1, None);
  88.         assert_eq!(dijkstra(&graph, &1), dists_1);
  89.     }
  91.     #[test]
  92.     fn tree_1() {
  93.         let mut graph = BTreeMap::new();
  94.         let mut dists = BTreeMap::new();
  95.         dists.insert(1, None);
  96.         for i in 1..100 {
  97.             add_edge(&mut graph, i, i * 2, i * 2);
  98.             add_edge(&mut graph, i, i * 2 + 1, i * 2 + 1);
  100.             match dists[&i] {
  101.                 Some((_, d)) => {
  102.                     dists.insert(i * 2, Some((i, d + i * 2)));
  103.                     dists.insert(i * 2 + 1, Some((i, d + i * 2 + 1)));
  104.                 }
  105.                 None => {
  106.                     dists.insert(i * 2, Some((i, i * 2)));
  107.                     dists.insert(i * 2 + 1, Some((i, i * 2 + 1)));
  108.                 }
  109.             }
  110.         }
  112.         assert_eq!(dijkstra(&graph, &1), dists);
  113.     }
  115.     #[test]
  116.     fn graph_1() {
  117.         let mut graph = BTreeMap::new();
  118.         add_edge(&mut graph, 'a', 'c', 12);
  119.         add_edge(&mut graph, 'a', 'd', 60);
  120.         add_edge(&mut graph, 'b', 'a', 10);
  121.         add_edge(&mut graph, 'c', 'b', 20);
  122.         add_edge(&mut graph, 'c', 'd', 32);
  123.         add_edge(&mut graph, 'e', 'a', 7);
  125.         let mut dists_a = BTreeMap::new();
  126.         dists_a.insert('a', None);
  127.         dists_a.insert('c', Some(('a', 12)));
  128.         dists_a.insert('d', Some(('c', 44)));
  129.         dists_a.insert('b', Some(('c', 32)));
  130.         assert_eq!(dijkstra(&graph, &'a'), dists_a);
  132.         let mut dists_b = BTreeMap::new();
  133.         dists_b.insert('b', None);
  134.         dists_b.insert('a', Some(('b', 10)));
  135.         dists_b.insert('c', Some(('a', 22)));
  136.         dists_b.insert('d', Some(('c', 54)));
  137.         assert_eq!(dijkstra(&graph, &'b'), dists_b);
  139.         let mut dists_c = BTreeMap::new();
  140.         dists_c.insert('c', None);
  141.         dists_c.insert('b', Some(('c', 20)));
  142.         dists_c.insert('d', Some(('c', 32)));
  143.         dists_c.insert('a', Some(('b', 30)));
  144.         assert_eq!(dijkstra(&graph, &'c'), dists_c);
  146.         let mut dists_d = BTreeMap::new();
  147.         dists_d.insert('d', None);
  148.         assert_eq!(dijkstra(&graph, &'d'), dists_d);
  150.         let mut dists_e = BTreeMap::new();
  151.         dists_e.insert('e', None);
  152.         dists_e.insert('a', Some(('e', 7)));
  153.         dists_e.insert('c', Some(('a', 19)));
  154.         dists_e.insert('d', Some(('c', 51)));
  155.         dists_e.insert('b', Some(('c', 39)));
  156.         assert_eq!(dijkstra(&graph, &'e'), dists_e);
  157.     }
  158. }