Implement A* search algorithm (TheAlgorithms#470)

Author: qsantos

src/graph/astar.rs

@@ -0,0 +1,262 @@
+use std::{
+    collections::{BTreeMap, BinaryHeap},
+    ops::Add,
+};
+
+use num_traits::Zero;
+
+type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+struct Candidate<V, E> {
+    estimated_weight: E,
+    real_weight: E,
+    state: V,
+}
+
+impl<V: Ord + Copy, E: Ord + Copy> PartialOrd for Candidate<V, E> {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        // Note the inverted order; we want nodes with lesser weight to have
+        // higher priority
+        other.estimated_weight.partial_cmp(&self.estimated_weight)
+    }
+}
+
+impl<V: Ord + Copy, E: Ord + Copy> Ord for Candidate<V, E> {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        // Note the inverted order; we want nodes with lesser weight to have
+        // higher priority
+        other.estimated_weight.cmp(&self.estimated_weight)
+    }
+}
+
+pub fn astar<V: Ord + Copy, E: Ord + Copy + Add<Output = E> + Zero>(
+    graph: &Graph<V, E>,
+    start: V,
+    target: V,
+    heuristic: impl Fn(V) -> E,
+) -> Option<(E, Vec<V>)> {
+    // traversal front
+    let mut queue = BinaryHeap::new();
+    // maps each node to its predecessor in the final path
+    let mut previous = BTreeMap::new();
+    // weights[v] is the accumulated weight from start to v
+    let mut weights = BTreeMap::new();
+    // initialize traversal
+    weights.insert(start, E::zero());
+    queue.push(Candidate {
+        estimated_weight: heuristic(start),
+        real_weight: E::zero(),
+        state: start,
+    });
+    while let Some(Candidate {
+        estimated_weight: _,
+        real_weight,
+        state: current,
+    }) = queue.pop()
+    {
+        if current == target {
+            break;
+        }
+        for (&next, &weight) in &graph[&current] {
+            let real_weight = real_weight + weight;
+            if weights
+                .get(&next)
+                .map(|&weight| real_weight < weight)
+                .unwrap_or(true)
+            {
+                // current allows us to reach next with lower weight (or at all)
+                // add next to the front
+                let estimated_weight = real_weight + heuristic(next);
+                weights.insert(next, real_weight);
+                queue.push(Candidate {
+                    estimated_weight,
+                    real_weight,
+                    state: next,
+                });
+                previous.insert(next, current);
+            }
+        }
+    }
+    let weight = if let Some(&weight) = weights.get(&target) {
+        weight
+    } else {
+        // we did not reach target from start
+        return None;
+    };
+    // build path in reverse
+    let mut current = target;
+    let mut path = vec![current];
+    while current != start {
+        let prev = previous
+            .get(&current)
+            .copied()
+            .expect("We reached the target, but are unable to reconsistute the path");
+        current = prev;
+        path.push(current);
+    }
+    path.reverse();
+    Some((weight, path))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{astar, Graph};
+    use num_traits::Zero;
+    use std::collections::BTreeMap;
+
+    // the null heuristic make A* equivalent to Dijkstra
+    fn null_heuristic<V, E: Zero>(_v: V) -> E {
+        E::zero()
+    }
+
+    fn add_edge<V: Ord + Copy, E: Ord>(graph: &mut Graph<V, E>, v1: V, v2: V, c: E) {
+        graph.entry(v1).or_insert_with(BTreeMap::new).insert(v2, c);
+        graph.entry(v2).or_insert_with(BTreeMap::new);
+    }
+
+    #[test]
+    fn single_vertex() {
+        let mut graph: Graph<usize, usize> = BTreeMap::new();
+        graph.insert(0, BTreeMap::new());
+
+        assert_eq!(astar(&graph, 0, 0, null_heuristic), Some((0, vec![0])));
+        assert_eq!(astar(&graph, 0, 1, null_heuristic), None);
+    }
+
+    #[test]
+    fn single_edge() {
+        let mut graph = BTreeMap::new();
+        add_edge(&mut graph, 0, 1, 2);
+
+        assert_eq!(astar(&graph, 0, 1, null_heuristic), Some((2, vec![0, 1])));
+        assert_eq!(astar(&graph, 1, 0, null_heuristic), None);
+    }
+
+    #[test]
+    fn graph_1() {
+        let mut graph = BTreeMap::new();
+        add_edge(&mut graph, 'a', 'c', 12);
+        add_edge(&mut graph, 'a', 'd', 60);
+        add_edge(&mut graph, 'b', 'a', 10);
+        add_edge(&mut graph, 'c', 'b', 20);
+        add_edge(&mut graph, 'c', 'd', 32);
+        add_edge(&mut graph, 'e', 'a', 7);
+
+        // from a
+        assert_eq!(
+            astar(&graph, 'a', 'a', null_heuristic),
+            Some((0, vec!['a']))
+        );
+        assert_eq!(
+            astar(&graph, 'a', 'b', null_heuristic),
+            Some((32, vec!['a', 'c', 'b']))
+        );
+        assert_eq!(
+            astar(&graph, 'a', 'c', null_heuristic),
+            Some((12, vec!['a', 'c']))
+        );
+        assert_eq!(
+            astar(&graph, 'a', 'd', null_heuristic),
+            Some((12 + 32, vec!['a', 'c', 'd']))
+        );
+        assert_eq!(astar(&graph, 'a', 'e', null_heuristic), None);
+
+        // from b
+        assert_eq!(
+            astar(&graph, 'b', 'a', null_heuristic),
+            Some((10, vec!['b', 'a']))
+        );
+        assert_eq!(
+            astar(&graph, 'b', 'b', null_heuristic),
+            Some((0, vec!['b']))
+        );
+        assert_eq!(
+            astar(&graph, 'b', 'c', null_heuristic),
+            Some((10 + 12, vec!['b', 'a', 'c']))
+        );
+        assert_eq!(
+            astar(&graph, 'b', 'd', null_heuristic),
+            Some((10 + 12 + 32, vec!['b', 'a', 'c', 'd']))
+        );
+        assert_eq!(astar(&graph, 'b', 'e', null_heuristic), None);
+
+        // from c
+        assert_eq!(
+            astar(&graph, 'c', 'a', null_heuristic),
+            Some((20 + 10, vec!['c', 'b', 'a']))
+        );
+        assert_eq!(
+            astar(&graph, 'c', 'b', null_heuristic),
+            Some((20, vec!['c', 'b']))
+        );
+        assert_eq!(
+            astar(&graph, 'c', 'c', null_heuristic),
+            Some((0, vec!['c']))
+        );
+        assert_eq!(
+            astar(&graph, 'c', 'd', null_heuristic),
+            Some((32, vec!['c', 'd']))
+        );
+        assert_eq!(astar(&graph, 'c', 'e', null_heuristic), None);
+
+        // from d
+        assert_eq!(astar(&graph, 'd', 'a', null_heuristic), None);
+        assert_eq!(astar(&graph, 'd', 'b', null_heuristic), None);
+        assert_eq!(astar(&graph, 'd', 'c', null_heuristic), None);
+        assert_eq!(
+            astar(&graph, 'd', 'd', null_heuristic),
+            Some((0, vec!['d']))
+        );
+        assert_eq!(astar(&graph, 'd', 'e', null_heuristic), None);
+
+        // from e
+        assert_eq!(
+            astar(&graph, 'e', 'a', null_heuristic),
+            Some((7, vec!['e', 'a']))
+        );
+        assert_eq!(
+            astar(&graph, 'e', 'b', null_heuristic),
+            Some((7 + 12 + 20, vec!['e', 'a', 'c', 'b']))
+        );
+        assert_eq!(
+            astar(&graph, 'e', 'c', null_heuristic),
+            Some((7 + 12, vec!['e', 'a', 'c']))
+        );
+        assert_eq!(
+            astar(&graph, 'e', 'd', null_heuristic),
+            Some((7 + 12 + 32, vec!['e', 'a', 'c', 'd']))
+        );
+        assert_eq!(
+            astar(&graph, 'e', 'e', null_heuristic),
+            Some((0, vec!['e']))
+        );
+    }
+
+    #[test]
+    fn test_heuristic() {
+        // make a grid
+        let mut graph = BTreeMap::new();
+        let rows = 100;
+        let cols = 100;
+        for row in 0..rows {
+            for col in 0..cols {
+                add_edge(&mut graph, (row, col), (row + 1, col), 1);
+                add_edge(&mut graph, (row, col), (row, col + 1), 1);
+                add_edge(&mut graph, (row, col), (row + 1, col + 1), 1);
+                add_edge(&mut graph, (row + 1, col), (row, col), 1);
+                add_edge(&mut graph, (row + 1, col + 1), (row, col), 1);
+            }
+        }
+
+        // Dijkstra would explore most of the 101 × 101 nodes
+        // the heuristic should allow exploring only about 200 nodes
+        let now = std::time::Instant::now();
+        let res = astar(&graph, (0, 0), (100, 90), |(i, j)| 100 - i + 90 - j);
+        assert!(now.elapsed() < std::time::Duration::from_millis(10));
+
+        let (weight, path) = res.unwrap();
+        assert_eq!(weight, 100);
+        assert_eq!(path.len(), 101);
+    }
+}

src/graph/dijkstra.rs

@@ -11,37 +11,37 @@ type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;
 // since the start has no predecessor but is reachable, map[start] will be None
 pub fn dijkstra<V: Ord + Copy, E: Ord + Copy + Add<Output = E>>(
     graph: &Graph<V, E>,
-    start: &V,
+    start: V,
 ) -> BTreeMap<V, Option<(V, E)>> {
     let mut ans = BTreeMap::new();
     let mut prio = BinaryHeap::new();
 
     // start is the special case that doesn't have a predecessor
-    ans.insert(*start, None);
+    ans.insert(start, None);
 
-    for (new, weight) in &graph[start] {
-        ans.insert(*new, Some((*start, *weight)));
-        prio.push(Reverse((*weight, new, start)));
+    for (new, weight) in &graph[&start] {
+        ans.insert(*new, Some((start, *weight)));
+        prio.push(Reverse((*weight, *new, start)));
     }
 
     while let Some(Reverse((dist_new, new, prev))) = prio.pop() {
-        match ans[new] {
+        match ans[&new] {
             // what we popped is what is in ans, we'll compute it
-            Some((p, d)) if p == *prev && d == dist_new => {}
+            Some((p, d)) if p == prev && d == dist_new => {}
             // otherwise it's not interesting
             _ => continue,
         }
 
-        for (next, weight) in &graph[new] {
+        for (next, weight) in &graph[&new] {
             match ans.get(next) {
                 // if ans[next] is a lower dist than the alternative one, we do nothing
                 Some(Some((_, dist_next))) if dist_new + *weight >= *dist_next => {}
                 // 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
                 Some(None) => {}
                 // the new path is shorter, either new was not in ans or it was farther
                 _ => {
-                    ans.insert(*next, Some((*new, *weight + dist_new)));
-                    prio.push(Reverse((*weight + dist_new, next, new)));
+                    ans.insert(*next, Some((new, *weight + dist_new)));
+                    prio.push(Reverse((*weight + dist_new, *next, new)));
                 }
             }
         }
@@ -68,7 +68,7 @@ mod tests {
         let mut dists = BTreeMap::new();
         dists.insert(0, None);
 
-        assert_eq!(dijkstra(&graph, &0), dists);
+        assert_eq!(dijkstra(&graph, 0), dists);
     }
 
     #[test]
@@ -80,12 +80,12 @@ mod tests {
         dists_0.insert(0, None);
         dists_0.insert(1, Some((0, 2)));
 
-        assert_eq!(dijkstra(&graph, &0), dists_0);
+        assert_eq!(dijkstra(&graph, 0), dists_0);
 
         let mut dists_1 = BTreeMap::new();
         dists_1.insert(1, None);
 
-        assert_eq!(dijkstra(&graph, &1), dists_1);
+        assert_eq!(dijkstra(&graph, 1), dists_1);
     }
 
     #[test]
@@ -109,7 +109,7 @@ mod tests {
             }
         }
 
-        assert_eq!(dijkstra(&graph, &1), dists);
+        assert_eq!(dijkstra(&graph, 1), dists);
     }
 
     #[test]
@@ -127,32 +127,32 @@ mod tests {
         dists_a.insert('c', Some(('a', 12)));
         dists_a.insert('d', Some(('c', 44)));
         dists_a.insert('b', Some(('c', 32)));
-        assert_eq!(dijkstra(&graph, &'a'), dists_a);
+        assert_eq!(dijkstra(&graph, 'a'), dists_a);
 
         let mut dists_b = BTreeMap::new();
         dists_b.insert('b', None);
         dists_b.insert('a', Some(('b', 10)));
         dists_b.insert('c', Some(('a', 22)));
         dists_b.insert('d', Some(('c', 54)));
-        assert_eq!(dijkstra(&graph, &'b'), dists_b);
+        assert_eq!(dijkstra(&graph, 'b'), dists_b);
 
         let mut dists_c = BTreeMap::new();
         dists_c.insert('c', None);
         dists_c.insert('b', Some(('c', 20)));
         dists_c.insert('d', Some(('c', 32)));
         dists_c.insert('a', Some(('b', 30)));
-        assert_eq!(dijkstra(&graph, &'c'), dists_c);
+        assert_eq!(dijkstra(&graph, 'c'), dists_c);
 
         let mut dists_d = BTreeMap::new();
         dists_d.insert('d', None);
-        assert_eq!(dijkstra(&graph, &'d'), dists_d);
+        assert_eq!(dijkstra(&graph, 'd'), dists_d);
 
         let mut dists_e = BTreeMap::new();
         dists_e.insert('e', None);
         dists_e.insert('a', Some(('e', 7)));
         dists_e.insert('c', Some(('a', 19)));
         dists_e.insert('d', Some(('c', 51)));
         dists_e.insert('b', Some(('c', 39)));
-        assert_eq!(dijkstra(&graph, &'e'), dists_e);
+        assert_eq!(dijkstra(&graph, 'e'), dists_e);
     }
 }

src/graph/mod.rs

@@ -1,3 +1,4 @@
+mod astar;
 mod bellman_ford;
 mod bipartite_matching;
 mod breadth_first_search;
@@ -17,6 +18,7 @@ mod prufer_code;
 mod strongly_connected_components;
 mod topological_sort;
 mod two_satisfiability;
+pub use self::astar::astar;
 pub use self::bellman_ford::bellman_ford;
 pub use self::bipartite_matching::BipartiteMatching;
 pub use self::breadth_first_search::breadth_first_search;