Stolen dijkstra time

src/day15.rs

@@ -1,9 +1,11 @@
-use std::cmp::max;
-use std::collections::{HashMap, HashSet};
+use std::cmp::{max, Reverse};
+use std::collections::{BinaryHeap, BTreeMap, HashMap};
 use crate::Day;
 
 type Coordinate = (usize, usize);
 
+type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;
+
 pub struct Day15(HashMap<Coordinate, u32>, usize, usize);
 
 impl Day for Day15 {
@@ -23,8 +25,10 @@ impl Day for Day15 {
     }
 
     fn part1(&self) -> anyhow::Result<String> {
-        let r = do_part_1(&self.0, (self.1, self.1));
-        Ok(format!("{}", r))
+        let g = map_to_graph(&self.0, (self.1, self.1));
+        let end = &(self.1, self.2);
+        let r = dijkstra(&g, &(0, 0), end).unwrap();
+        Ok(format!("{}", r.1))
     }
 
     fn part2(&self) -> anyhow::Result<String> {
@@ -37,7 +41,6 @@ impl Day for Day15 {
                 for ry in 0..5 {
                     for rx in 0..5 {
                         let coord = (rx * len + x, ry * len + y);
-                        // println!("{:?}", coord);
                         let e = m.entry(coord).or_insert(0);
                         *e = orig_value + rx as u32 + ry as u32;
                         if *e > 9 {
@@ -47,7 +50,11 @@ impl Day for Day15 {
                 }
             }
         }
-        Ok(format!("{}", do_part_1(&m, (len * 5 - 1, len * 5 - 1))))
+
+        let target = (len * 5 - 1, len * 5 - 1);
+        let g = map_to_graph(&m, target);
+        let r = dijkstra(&g, &(0, 0), &target).unwrap();
+        Ok(format!("{}", r.1))
     }
 }
 
@@ -60,33 +67,63 @@ fn neighbours((x, y): Coordinate, (tx, ty): Coordinate) -> impl Iterator<Item=Co
     ].into_iter().flatten()
 }
 
+pub fn dijkstra(
+    graph: &Graph<Coordinate, u32>,
+    start: &Coordinate,
+    end: &Coordinate,
+) -> Option<(Coordinate, u32)> {
+    let mut ans = BTreeMap::new();
+    let mut prio = BinaryHeap::new();
 
-fn do_part_1(map: &HashMap<Coordinate, u32>, target: Coordinate) -> u32 {
-    let mut distance = HashMap::new();
-    let mut seen = HashSet::new();
-    let mut heap = vec![(0, 0)];
-    distance.insert((0, 0), 0);
+    // start is the special case that doesn't have a predecessor
+    ans.insert(*start, None);
 
-    while !heap.is_empty() {
-        heap.sort_unstable_by(|a, b| distance.get(b).unwrap().cmp(distance.get(a).unwrap()));
+    for (new, weight) in &graph[start] {
+        ans.insert(*new, Some((*start, *weight)));
+        prio.push(Reverse((*weight, new, start)));
+    }
 
-        let nearest = heap.pop().unwrap();
-        if nearest == target {
-            break;
+    while let Some(Reverse((dist_new, new, prev))) = prio.pop() {
+        match ans[new] {
+            // what we popped is what is in ans, we'll compute it
+            Some((p, d)) if p == *prev && d == dist_new => {}
+            // otherwise it's not interesting
+            _ => continue,
         }
-        seen.insert(nearest);
-        for neighbour in neighbours(nearest, target) {
-            if !seen.contains(&neighbour) && !heap.contains(&neighbour) {
-                heap.push(neighbour);
-                let risk = distance.get(&nearest).unwrap() + map.get(&neighbour).unwrap();
-                if risk < *distance.get(&neighbour).unwrap_or(&u32::MAX) {
-                    distance.insert(neighbour, risk);
+
+        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)));
                 }
             }
         }
     }
 
-    *distance.get(&target).unwrap()
+    ans.get(end).unwrap().clone()
+}
+
+fn add_edge(graph: &mut Graph<Coordinate, u32>, source: Coordinate, target: Coordinate, value: u32) {
+    graph.entry(source).or_insert_with(BTreeMap::new).insert(target, value);
+    graph.entry(target).or_insert_with(BTreeMap::new);
+}
+
+fn map_to_graph(map: &HashMap<Coordinate, u32>, end: Coordinate) -> Graph<Coordinate, u32> {
+    let mut g = BTreeMap::new();
+    for y in 0..=end.1 {
+        for x in 0..=end.0 {
+            for neighbour in neighbours((x, y), end) {
+                add_edge(&mut g, (x, y), neighbour, *map.get(&neighbour).unwrap())
+            }
+        }
+    }
+    g
 }
 
 #[cfg(test)]
@@ -120,17 +157,17 @@ mod tests {
     }
 
     // Woo slow and i again don't feel like fixing it :p
-    // #[test]
-    // fn part1_real() -> anyhow::Result<()> {
-    //     let d = Day15::init(crate::load_input("15")?)?;
-    //     assert_eq!("388", d.part1()?);
-    //     Ok(())
-    // }
-    //
-    // #[test]
-    // fn part2_real() -> anyhow::Result<()> {
-    //     let d = Day15::init(crate::load_input("15")?)?;
-    //     assert_eq!("2819", d.part2()?);
-    //     Ok(())
-    // }
+    #[test]
+    fn part1_real() -> anyhow::Result<()> {
+        let d = Day15::init(crate::load_input("15")?)?;
+        assert_eq!("388", d.part1()?);
+        Ok(())
+    }
+
+    #[test]
+    fn part2_real() -> anyhow::Result<()> {
+        let d = Day15::init(crate::load_input("15")?)?;
+        assert_eq!("2819", d.part2()?);
+        Ok(())
+    }
 }