feat: impl day 17

3 files changed, 113 insertions, 91 deletions

diff --git a/src/day_17.rs b/src/day_17.rs
index 0f90905..0ba3be8 100644
--- a/src/day_17.rs
+++ b/src/day_17.rs
@@ -1,4 +1,7 @@
-use std::collections::{hash_map::Entry, BinaryHeap, HashMap, HashSet, VecDeque};
+use std::{
+    cmp::Reverse,
+    collections::{hash_map::Entry, BinaryHeap, HashMap},
+};
 
 use anyhow::Context;
 
@@ -20,32 +23,19 @@ impl Solution for Day17 {
     fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
         let grid = parse_grid(input).context("Failed to parse grid")?;
 
-        let path = astar(&grid, (0, 0), (grid.len() - 1, grid[0].len() - 1))
-            .context("Failed to find path")?;
+        astar::<0, 3>(&grid, (0, 0), (grid.len() - 1, grid[0].len() - 1))
+            .context("Failed to find path")
 
-        let mut visited = HashSet::new();
-        for (p, cost) in path {
-            visited.insert(p);
-            println!("{p:?}: {cost}");
-        }
-
-        for (row, v) in grid.iter().enumerate() {
-            let mut output = String::new();
-            for (col, _) in v.iter().enumerate() {
-                if visited.contains(&(row, col)) {
-                    output.push('#');
-                } else {
-                    output.push('.');
-                }
-            }
-            println!("{output}")
-        }
-
-        todo!()
+        // 755 ==
     }
 
     fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
-        todo!()
+        let grid = parse_grid(input).context("Failed to parse grid")?;
+
+        astar::<4, 10>(&grid, (0, 0), (grid.len() - 1, grid[0].len() - 1))
+            .context("Failed to find path")
+        // 879 <
+        // 881 ==
     }
 }
 
@@ -60,17 +50,31 @@ fn parse_grid(input: &str) -> Option<Grid> {
         .try_collect::<Grid>()
 }
 
-fn astar(grid: &Grid, start: Position, goal: Position) -> Option<Vec<(Position, usize)>> {
-    let mut queue = BinaryHeap::from([Node::start(start)]);
-    let mut cache = HashMap::from([(start, usize::MAX)]);
+fn astar<const M: u8, const N: u8>(grid: &Grid, start: Position, goal: Position) -> Option<usize> {
+    let mut queue = BinaryHeap::from([Reverse(Node::<M, N>::start(start))]);
+    let mut cache = HashMap::from([((start, vec![]), usize::MAX)]);
+
+    while let Some(Reverse(node)) = queue.pop() {
+        for (edge, position) in node.successors(grid) {
+            let mut next = node.clone();
+            next.position = position;
+            next.path.push(edge);
+            next.cost += grid[position.0][position.1];
+            next.estimated = next.cost + next.dist(goal);
 
-    while let Some(node) = queue.pop() {
-        for next in node.successors(grid, goal) {
-            if next.position == goal {
-                return Some(rebuild_path(cache, next));
+            if next.goal(goal) {
+                return Some(next.cost);
             }
 
-            match cache.entry(next.position) {
+            let trail = next
+                .path
+                .iter()
+                .copied()
+                .rev()
+                .take(N as usize)
+                .rev()
+                .collect::<Vec<_>>();
+            match cache.entry((next.position, trail)) {
                 Entry::Vacant(e) => {
                     e.insert(next.cost);
                 }
@@ -83,36 +87,25 @@ fn astar(grid: &Grid, start: Position, goal: Position) -> Option<Vec<(Position,
                 }
             }
 
-            queue.push(next)
+            queue.push(Reverse(next))
         }
     }
 
     None
 }
 
-fn rebuild_path(cache: HashMap<Position, usize>, node: Node) -> Vec<(Position, usize)> {
-    let mut path = node
-        .parents
-        .into_iter()
-        .map(|p| (p, *cache.get(&p).unwrap()))
-        .collect::<Vec<_>>();
-    path.push((node.position, node.cost));
-    path
-}
-
 type Grid = Vec<Vec<usize>>;
 type Position = (usize, usize);
 
-#[derive(Debug, Default, Clone, PartialEq, Eq)]
-struct Node {
+#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub struct Node<const M: u8, const N: u8> {
     estimated: usize,
     cost: usize,
     position: Position,
-    parents: Vec<Position>,
-    trail: VecDeque<Edge>,
+    path: Vec<Edge>,
 }
 
-impl Node {
+impl<const M: u8, const N: u8> Node<M, N> {
     fn start(position: Position) -> Self {
         Node {
             position,
@@ -121,62 +114,62 @@ impl Node {
         }
     }
 
-    fn successors(&self, grid: &Grid, goal: Position) -> Vec<Self> {
-        Edge::ALL
+    fn goal(&self, position: Position) -> bool {
+        self.position == position
+            && self
+                .path
+                .iter()
+                .rev()
+                .take(M as usize)
+                .map_windows(|[a, b]| a == b)
+                .all(|b| b)
+    }
+
+    fn successors(&self, grid: &Grid) -> Vec<(Edge, Position)> {
+        let Some(last) = self.path.last() else {
+            return Edge::ALL
+                .into_iter()
+                .filter_map(|edge| self.traverse_edge(grid, edge))
+                .collect::<Vec<_>>();
+        };
+
+        last.edges()
             .into_iter()
-            .filter(|e| self.trail.len() < 3 || !self.trail.iter().all(|t| e == t))
-            .inspect(|e| print!(" {e:?}"))
-            .filter_map(|e| {
-                let (row, col) = self.position;
-
-                let position = match e {
-                    Edge::North if row < grid.len() - 1 => (row.checked_add(1)?, col),
-                    Edge::East if col < grid[row].len() - 1 => (row, col.checked_add(1)?),
-                    Edge::South => (row.checked_sub(1)?, col),
-                    Edge::West => (row, col.checked_sub(1)?),
-                    _ => return None,
-                };
-
-                let mut next = self.clone();
-                next.position = position;
-                next.cost += grid[position.0][position.1];
-                next.estimated = next.cost + next.dist(goal);
-                next.parents.push(self.position);
-
-                if next.trail.len() >= 3 {
-                    next.trail.pop_front();
+            .filter(|edge| {
+                let seq = self.path.iter().rev().take_while(|e| *e == last).count();
+                if edge == last {
+                    seq < N as usize
+                } else {
+                    seq >= M as usize
                 }
-                next.trail.push_back(e);
-
-                Some(next)
             })
+            .filter_map(|edge| self.traverse_edge(grid, edge))
             .collect()
     }
 
-    fn dist(&self, p2: Position) -> usize {
-        self.position.0.abs_diff(p2.0) + self.position.1.abs_diff(p2.1)
+    fn traverse_edge(&self, grid: &Grid, edge: Edge) -> Option<(Edge, Position)> {
+        let (row, col) = self.position;
+
+        Some((
+            edge,
+            match edge {
+                Edge::North => (row.checked_sub(1)?, col),
+                Edge::East if col < grid[row].len().saturating_sub(1) => (row, col.checked_add(1)?),
+                Edge::South if row < grid.len().saturating_sub(1) => (row.checked_add(1)?, col),
+                Edge::West => (row, col.checked_sub(1)?),
+                _ => return None,
+            },
+        ))
     }
-}
 
-impl PartialOrd for Node {
-    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl Ord for Node {
-    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        match self.estimated.cmp(&other.estimated) {
-            std::cmp::Ordering::Equal => self.cost.cmp(&other.cost),
-            c => c,
-        }
+    fn dist(&self, p2: Position) -> usize {
+        self.position.0.abs_diff(p2.0) + self.position.1.abs_diff(p2.1)
     }
 }
 
-#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 enum Edge {
     North,
-    #[default]
     East,
     South,
     West,
@@ -184,6 +177,15 @@ enum Edge {
 
 impl Edge {
     const ALL: [Edge; 4] = [Edge::North, Edge::East, Edge::South, Edge::West];
+
+    fn edges(&self) -> [Edge; 3] {
+        match self {
+            Edge::North => [Edge::North, Edge::East, Edge::West],
+            Edge::East => [Edge::North, Edge::East, Edge::South],
+            Edge::South => [Edge::East, Self::South, Edge::West],
+            Edge::West => [Edge::North, Edge::South, Edge::West],
+        }
+    }
 }
 
 #[cfg(test)]
@@ -206,8 +208,26 @@ mod tests {
         4322674655533
     "#};
 
+    const INPUT2: &str = indoc::indoc! {r#"
+        111111111111
+        999999999991
+        999999999991
+        999999999991
+        999999999991
+    "#};
+
     #[test]
     fn test_part_1() -> anyhow::Result<()> {
         Ok(assert_eq!(102, Day17::part_1(INPUT)?))
     }
+
+    #[test]
+    fn test_part_2_1() -> anyhow::Result<()> {
+        Ok(assert_eq!(94, Day17::part_2(INPUT)?))
+    }
+
+    #[test]
+    fn test_part_2_2() -> anyhow::Result<()> {
+        Ok(assert_eq!(71, Day17::part_2(INPUT2)?))
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 47562cb..51f906c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,6 @@
 #![feature(
     iterator_try_collect,
+    iter_map_windows,
     iter_array_chunks,
     array_windows,
     iter_intersperse,
diff --git a/src/main.rs b/src/main.rs
index 5af80db..33333c1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,7 +1,7 @@
 use aoc_2023::{
     day_01::Day01, day_02::Day02, day_03::Day03, day_04::Day04, day_05::Day05, day_06::Day06,
     day_07::Day07, day_08::Day08, day_09::Day09, day_10::Day10, day_11::Day11, day_12::Day12,
-    day_13::Day13, day_14::Day14, day_15::Day15, day_16::Day16, Solution,
+    day_13::Day13, day_14::Day14, day_15::Day15, day_16::Day16, day_17::Day17, Solution,
 };
 
 fn main() -> anyhow::Result<()> {
@@ -21,6 +21,7 @@ fn main() -> anyhow::Result<()> {
     Day14::solve()?;
     Day15::solve()?;
     Day16::solve()?;
+    Day17::solve()?;
 
     Ok(())
 }