最短路径-Bellman Ford

  1. use std::collections::BTreeMap;
  2. use std::ops::Add;
  4. use std::ops::Neg;
  6. type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;
  8. // performs the Bellman-Ford algorithm on the given graph from the given start
  9. // the graph is an undirected graph
  10. //
  11. // if there is a negative weighted loop it returns None
  12. // else it returns a map that for each reachable vertex associates the distance and the predecessor
  13. // since the start has no predecessor but is reachable, map[start] will be None
  14. pub fn bellman_ford<
  15.     V: Ord + Copy,
  16.     E: Ord + Copy + Add<Output = E> + Neg<Output = E> + std::ops::Sub<Output = E>,
  17. >(
  18.     graph: &Graph<V, E>,
  19.     start: &V,
  20. ) -> Option<BTreeMap<V, Option<(V, E)>>> {
  21.     let mut ans: BTreeMap<V, Option<(V, E)>> = BTreeMap::new();
  23.     ans.insert(*start, None);
  25.     for _ in 1..(graph.len()) {
  26.         for (u, edges) in graph {
  27.             let dist_u = match ans.get(u) {
  28.                 Some(Some((_, d))) => Some(*d),
  29.                 Some(None) => None,
  30.                 None => continue,
  31.             };
  33.             for (v, d) in edges {
  34.                 match ans.get(v) {
  35.                     Some(Some((_, dist)))
  36.                         // if this is a longer path, do nothing
  37.                         if match dist_u {
  38.                             Some(dist_u) => dist_u + *d >= *dist,
  39.                             None => d >= dist,
  40.                         } => {}
  41.                     Some(None) => {
  42.                         match dist_u {
  43.                             // if dist_u + d < 0 there is a negative loop going by start
  44.                             // else it's just a longer path
  45.                             Some(dist_u) if dist_u >= -*d => {}
  46.                             // negative self edge or negative loop
  47.                             _ => {
  48.                                 if *d > *d + *d {
  49.                                     return None;
  50.                                 }
  51.                             }
  52.                         };
  53.                     }
  54.                     // it's a shorter path: either dist_v was infinite or it was longer than dist_u + d
  55.                     _ => {
  56.                         ans.insert(
  57.                             *v,
  58.                             Some((
  59.                                 *u,
  60.                                 match dist_u {
  61.                                     Some(dist) => dist + *d,
  62.                                     None => *d,
  63.                                 },
  64.                             )),
  65.                         );
  66.                     }
  67.                 }
  68.             }
  69.         }
  70.     }
  72.     for (u, edges) in graph {
  73.         for (v, d) in edges {
  74.             match (ans.get(u), ans.get(v)) {
  75.                 (Some(None), Some(None)) if *d > *d + *d => return None,
  76.                 (Some(None), Some(Some((_, dv)))) if d < dv => return None,
  77.                 (Some(Some((_, du))), Some(None)) if *du < -*d => return None,
  78.                 (Some(Some((_, du))), Some(Some((_, dv)))) if *du + *d < *dv => return None,
  79.                 (_, _) => {}
  80.             }
  81.         }
  82.     }
  84.     Some(ans)
  85. }
  87. #[cfg(test)]
  88. mod tests {
  89.     use super::{bellman_ford, Graph};
  90.     use std::collections::BTreeMap;
  92.     fn add_edge<V: Ord + Copy, E: Ord>(graph: &mut Graph<V, E>, v1: V, v2: V, c: E) {
  93.         graph.entry(v1).or_insert_with(BTreeMap::new).insert(v2, c);
  94.         graph.entry(v2).or_insert_with(BTreeMap::new);
  95.     }
  97.     #[test]
  98.     fn single_vertex() {
  99.         let mut graph: Graph<isize, isize> = BTreeMap::new();
  100.         graph.insert(0, BTreeMap::new());
  102.         let mut dists = BTreeMap::new();
  103.         dists.insert(0, None);
  105.         assert_eq!(bellman_ford(&graph, &0), Some(dists));
  106.     }
  108.     #[test]
  109.     fn single_edge() {
  110.         let mut graph = BTreeMap::new();
  111.         add_edge(&mut graph, 0, 1, 2);
  113.         let mut dists_0 = BTreeMap::new();
  114.         dists_0.insert(0, None);
  115.         dists_0.insert(1, Some((0, 2)));
  117.         assert_eq!(bellman_ford(&graph, &0), Some(dists_0));
  119.         let mut dists_1 = BTreeMap::new();
  120.         dists_1.insert(1, None);
  122.         assert_eq!(bellman_ford(&graph, &1), Some(dists_1));
  123.     }
  125.     #[test]
  126.     fn tree_1() {
  127.         let mut graph = BTreeMap::new();
  128.         let mut dists = BTreeMap::new();
  129.         dists.insert(1, None);
  130.         for i in 1..100 {
  131.             add_edge(&mut graph, i, i * 2, i * 2);
  132.             add_edge(&mut graph, i, i * 2 + 1, i * 2 + 1);
  134.             match dists[&i] {
  135.                 Some((_, d)) => {
  136.                     dists.insert(i * 2, Some((i, d + i * 2)));
  137.                     dists.insert(i * 2 + 1, Some((i, d + i * 2 + 1)));
  138.                 }
  139.                 None => {
  140.                     dists.insert(i * 2, Some((i, i * 2)));
  141.                     dists.insert(i * 2 + 1, Some((i, i * 2 + 1)));
  142.                 }
  143.             }
  144.         }
  146.         assert_eq!(bellman_ford(&graph, &1), Some(dists));
  147.     }
  149.     #[test]
  150.     fn graph_1() {
  151.         let mut graph = BTreeMap::new();
  152.         add_edge(&mut graph, 'a', 'c', 12);
  153.         add_edge(&mut graph, 'a', 'd', 60);
  154.         add_edge(&mut graph, 'b', 'a', 10);
  155.         add_edge(&mut graph, 'c', 'b', 20);
  156.         add_edge(&mut graph, 'c', 'd', 32);
  157.         add_edge(&mut graph, 'e', 'a', 7);
  159.         let mut dists_a = BTreeMap::new();
  160.         dists_a.insert('a', None);
  161.         dists_a.insert('c', Some(('a', 12)));
  162.         dists_a.insert('d', Some(('c', 44)));
  163.         dists_a.insert('b', Some(('c', 32)));
  164.         assert_eq!(bellman_ford(&graph, &'a'), Some(dists_a));
  166.         let mut dists_b = BTreeMap::new();
  167.         dists_b.insert('b', None);
  168.         dists_b.insert('a', Some(('b', 10)));
  169.         dists_b.insert('c', Some(('a', 22)));
  170.         dists_b.insert('d', Some(('c', 54)));
  171.         assert_eq!(bellman_ford(&graph, &'b'), Some(dists_b));
  173.         let mut dists_c = BTreeMap::new();
  174.         dists_c.insert('c', None);
  175.         dists_c.insert('b', Some(('c', 20)));
  176.         dists_c.insert('d', Some(('c', 32)));
  177.         dists_c.insert('a', Some(('b', 30)));
  178.         assert_eq!(bellman_ford(&graph, &'c'), Some(dists_c));
  180.         let mut dists_d = BTreeMap::new();
  181.         dists_d.insert('d', None);
  182.         assert_eq!(bellman_ford(&graph, &'d'), Some(dists_d));
  184.         let mut dists_e = BTreeMap::new();
  185.         dists_e.insert('e', None);
  186.         dists_e.insert('a', Some(('e', 7)));
  187.         dists_e.insert('c', Some(('a', 19)));
  188.         dists_e.insert('d', Some(('c', 51)));
  189.         dists_e.insert('b', Some(('c', 39)));
  190.         assert_eq!(bellman_ford(&graph, &'e'), Some(dists_e));
  191.     }
  193.     #[test]
  194.     fn graph_2() {
  195.         let mut graph = BTreeMap::new();
  196.         add_edge(&mut graph, 0, 1, 6);
  197.         add_edge(&mut graph, 0, 3, 7);
  198.         add_edge(&mut graph, 1, 2, 5);
  199.         add_edge(&mut graph, 1, 3, 8);
  200.         add_edge(&mut graph, 1, 4, -4);
  201.         add_edge(&mut graph, 2, 1, -2);
  202.         add_edge(&mut graph, 3, 2, -3);
  203.         add_edge(&mut graph, 3, 4, 9);
  204.         add_edge(&mut graph, 4, 0, 3);
  205.         add_edge(&mut graph, 4, 2, 7);
  207.         let mut dists_0 = BTreeMap::new();
  208.         dists_0.insert(0, None);
  209.         dists_0.insert(1, Some((2, 2)));
  210.         dists_0.insert(2, Some((3, 4)));
  211.         dists_0.insert(3, Some((0, 7)));
  212.         dists_0.insert(4, Some((1, -2)));
  213.         assert_eq!(bellman_ford(&graph, &0), Some(dists_0));
  215.         let mut dists_1 = BTreeMap::new();
  216.         dists_1.insert(0, Some((4, -1)));
  217.         dists_1.insert(1, None);
  218.         dists_1.insert(2, Some((4, 3)));
  219.         dists_1.insert(3, Some((0, 6)));
  220.         dists_1.insert(4, Some((1, -4)));
  221.         assert_eq!(bellman_ford(&graph, &1), Some(dists_1));
  223.         let mut dists_2 = BTreeMap::new();
  224.         dists_2.insert(0, Some((4, -3)));
  225.         dists_2.insert(1, Some((2, -2)));
  226.         dists_2.insert(2, None);
  227.         dists_2.insert(3, Some((0, 4)));
  228.         dists_2.insert(4, Some((1, -6)));
  229.         assert_eq!(bellman_ford(&graph, &2), Some(dists_2));
  231.         let mut dists_3 = BTreeMap::new();
  232.         dists_3.insert(0, Some((4, -6)));
  233.         dists_3.insert(1, Some((2, -5)));
  234.         dists_3.insert(2, Some((3, -3)));
  235.         dists_3.insert(3, None);
  236.         dists_3.insert(4, Some((1, -9)));
  237.         assert_eq!(bellman_ford(&graph, &3), Some(dists_3));
  239.         let mut dists_4 = BTreeMap::new();
  240.         dists_4.insert(0, Some((4, 3)));
  241.         dists_4.insert(1, Some((2, 5)));
  242.         dists_4.insert(2, Some((4, 7)));
  243.         dists_4.insert(3, Some((0, 10)));
  244.         dists_4.insert(4, None);
  245.         assert_eq!(bellman_ford(&graph, &4), Some(dists_4));
  246.     }
  248.     #[test]
  249.     fn graph_with_negative_loop() {
  250.         let mut graph = BTreeMap::new();
  251.         add_edge(&mut graph, 0, 1, 6);
  252.         add_edge(&mut graph, 0, 3, 7);
  253.         add_edge(&mut graph, 1, 2, 5);
  254.         add_edge(&mut graph, 1, 3, 8);
  255.         add_edge(&mut graph, 1, 4, -4);
  256.         add_edge(&mut graph, 2, 1, -4);
  257.         add_edge(&mut graph, 3, 2, -3);
  258.         add_edge(&mut graph, 3, 4, 9);
  259.         add_edge(&mut graph, 4, 0, 3);
  260.         add_edge(&mut graph, 4, 2, 7);
  262.         assert_eq!(bellman_ford(&graph, &0), None);
  263.         assert_eq!(bellman_ford(&graph, &1), None);
  264.         assert_eq!(bellman_ford(&graph, &2), None);
  265.         assert_eq!(bellman_ford(&graph, &3), None);
  266.         assert_eq!(bellman_ford(&graph, &4), None);
  267.     }
  268. }