day22 2nd part

src/tasks/day22.rs

@@ -1,18 +1,22 @@
 use crate::tasks::day22::Equipment::*;
+use core::cmp::Ordering;
+use std::collections::BinaryHeap;
+use std::collections::HashMap;
 
-pub fn task1() {
+pub fn both() {
     let cave = Cave::create(3879, Node(8, 713, Torch));
-
     println!("Sum of erosion indexes: {}", cave.erosion_sum());
+    println!("Shortest path length: {}", cave.shortest_path_length());
 }
 
-#[derive(PartialEq)]
+#[derive(PartialEq, Clone, Copy, Debug, Eq, Hash)]
 enum Equipment {
     Torch,
     Climbing,
     Neither,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 struct Node(usize, usize, Equipment);
 
 struct Cave {
@@ -25,13 +29,15 @@ impl Cave {
         if target.2 != Torch {
             panic!("A valid target point needs the torch equipped");
         }
-        let mut map: Vec<Vec<usize>> = Vec::with_capacity(target.0 + 1);
-        let mut inner_vec = Vec::with_capacity(target.1 + 1);
-        inner_vec.resize(target.1 + 1, 0);
-        map.resize(target.0 + 1, inner_vec);
+        let width = 5000;
+        let height = 5000;
+        let mut map: Vec<Vec<usize>> = Vec::with_capacity(width);
+        let mut inner_vec = Vec::with_capacity(height);
+        inner_vec.resize(height, 0);
+        map.resize(width, inner_vec);
 
-        for x in 0..=target.0 {
-            for y in 0..=target.1 {
+        for x in 0..width {
+            for y in 0..height {
                 let geo_index = match (x, y) {
                     (0, 0) => 0,
                     _ if target.0 == x && target.1 == y => 0,
@@ -60,4 +66,156 @@ impl Cave {
     fn field_type(&self, x: usize, y: usize) -> usize {
         self.map[x][y] % 3
     }
+
+    fn shortest_path_length(&self) -> usize {
+        let start = Node(0, 0, Torch);
+        let mut open: BinaryHeap<State> = BinaryHeap::new();
+        let mut distances: HashMap<Node, usize> = HashMap::new();
+
+        distances.insert(start, 0);
+        open.push(State {
+            cost: 0,
+            position: start,
+        });
+
+        // Examine the frontier with lower cost nodes first (min-heap)
+        while let Some(State { cost, position }) = open.pop() {
+            // Alternatively we could have continued to find all shortest paths
+            if position == self.target {
+                return cost;
+            }
+
+            // Important as we may have already found a better way
+            if cost > *distances.entry(position).or_insert(usize::max_value()) {
+                continue;
+            }
+
+            // For each node we can reach, see if we can find a way with
+            // a lower cost going through this node
+            for edge in self.neighbors(position) {
+                let next = State {
+                    cost: cost + edge.cost,
+                    position: edge.node,
+                };
+
+                // If so, add it to the frontier and continue
+                let current_distance = distances.entry(edge.node).or_insert(usize::max_value());
+                if next.cost < *current_distance {
+                    open.push(next);
+                    // Relaxation, we have now found a better way
+                    *current_distance = next.cost;
+                }
+            }
+        }
+
+        unreachable!("There is always a path");
+    }
+
+    fn neighbors(&self, position: Node) -> Vec<Edge> {
+        let mut result = Vec::new();
+        // add all variants of the current position
+        result.push(Edge {
+            cost: 7,
+            node: self.other_node_for_region(position),
+        });
+
+        // for any neighbor position: if it allows the same equipment and is within bounds: add it
+        [
+            (position.0 as i32 - 1, position.1 as i32),
+            (position.0 as i32 + 1, position.1 as i32),
+            (position.0 as i32, position.1 as i32 - 1),
+            (position.0 as i32, position.1 as i32 + 1),
+        ]
+        .into_iter()
+        .filter_map(|(x, y)| {
+            if *x >= 0 && *y >= 0 {
+                Some(Node(*x as usize, *y as usize, position.2))
+            } else {
+                None
+            }
+        })
+        .for_each(|node| {
+            if self.equipment_allowed_for_region(node.0, node.1, node.2) {
+                result.push(Edge {
+                    cost: 1,
+                    node: node,
+                })
+            }
+        });
+
+        result
+    }
+
+    fn equipment_allowed_for_region(&self, x: usize, y: usize, equipment: Equipment) -> bool {
+        let field_type = self.field_type(x, y);
+        match field_type {
+            // rocky
+            0 => equipment == Torch || equipment == Climbing,
+            // wet
+            1 => equipment == Climbing || equipment == Neither,
+            //narrow
+            2 => equipment == Torch || equipment == Neither,
+            _ => panic!("not a valid type!"),
+        }
+    }
+
+    fn other_node_for_region(&self, position: Node) -> Node {
+        let field_type = self.field_type(position.0, position.1);
+        Node(
+            position.0,
+            position.1,
+            match field_type {
+                0 => match position.2 {
+                    Climbing => Torch,
+                    Torch => Climbing,
+                    _ => panic!(),
+                },
+                1 => match position.2 {
+                    Climbing => Neither,
+                    Neither => Climbing,
+                    _ => panic!(),
+                },
+                2 => match position.2 {
+                    Torch => Neither,
+                    Neither => Torch,
+                    _ => panic!(),
+                },
+                _ => panic!("not a valid type"),
+            },
+        )
+    }
+}
+
+#[derive(Debug)]
+struct Edge {
+    cost: usize,
+    node: Node,
+}
+
+#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+struct State {
+    cost: usize,
+    position: Node,
+}
+
+// The priority queue depends on `Ord`.
+// Explicitly implement the trait so the queue becomes a min-heap
+// instead of a max-heap.
+impl Ord for State {
+    fn cmp(&self, other: &State) -> Ordering {
+        // Notice that the we flip the ordering on costs.
+        // In case of a tie we compare positions - this step is necessary
+        // to make implementations of `PartialEq` and `Ord` consistent.
+        other
+            .cost
+            .cmp(&self.cost)
+            .then_with(|| self.position.0.cmp(&other.position.0))
+    }
+}
+
+// `PartialOrd` needs to be implemented as well.
+impl PartialOrd for State {
+    fn partial_cmp(&self, other: &State) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
 }