solution: day 17

data/examples/17.txt

@@ -0,0 +1,13 @@
+2413432311323
+3215453535623
+3255245654254
+3446585845452
+4546657867536
+1438598798454
+4457876987766
+3637877979653
+4654967986887
+4564679986453
+1224686865563
+2546548887735
+4322674655533

src/bin/17.rs

@@ -0,0 +1,192 @@
+use std::{
+    collections::{BinaryHeap, HashMap},
+    ops::Neg,
+};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+enum Direction {
+    Up,
+    Down,
+    Left,
+    Right,
+}
+
+impl Direction {
+    const ALL: [Direction; 4] = [
+        Direction::Up,
+        Direction::Down,
+        Direction::Left,
+        Direction::Right,
+    ];
+}
+
+impl Neg for Direction {
+    type Output = Self;
+    fn neg(self) -> Self::Output {
+        use Direction::*;
+        match self {
+            Up => Down,
+            Down => Up,
+            Left => Right,
+            Right => Left,
+        }
+    }
+}
+
+#[derive(Debug, PartialEq, Eq)]
+struct State {
+    heat: usize,
+    position: (usize, usize),
+    direction: Direction,
+    steps: usize,
+}
+
+impl PartialOrd for State {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for State {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        other
+            .heat
+            .cmp(&self.heat)
+            .then_with(|| self.position.cmp(&other.position))
+    }
+}
+
+fn find_path(
+    grid: &[Vec<usize>],
+    start: (usize, usize),
+    target: (usize, usize),
+    min_steps: usize,
+    max_steps: usize,
+) -> Option<usize> {
+    let valid_indexing = |y: usize, x: usize| grid.get(y).and_then(|row| row.get(x)).is_some();
+
+    let mut heap = BinaryHeap::new();
+    heap.push(State {
+        heat: 0,
+        position: (0, 0),
+        direction: Direction::Right,
+        steps: 0,
+    });
+
+    let mut heatmap = HashMap::new();
+    heatmap.insert((start, Direction::Down, 0), 0);
+    heatmap.insert((start, Direction::Up, 0), 0);
+
+    while let Some(State {
+        heat,
+        position,
+        direction,
+        steps,
+    }) = heap.pop()
+    {
+        if position == target {
+            return Some(heat);
+        }
+
+        if *heatmap
+            .get(&(position, direction, steps))
+            .unwrap_or(&usize::MAX)
+            < heat
+        {
+            continue;
+        }
+
+        let (y, x) = position;
+        for d in Direction::ALL.iter().filter(|&x| *x != -direction) {
+            if steps < min_steps && *d != direction {
+                continue;
+            }
+
+            let (ny, nx) = match d {
+                Direction::Up => (y.wrapping_sub(1), x),
+                Direction::Down => (y + 1, x),
+                Direction::Left => (y, x.wrapping_sub(1)),
+                Direction::Right => (y, x + 1),
+            };
+
+            if !valid_indexing(ny, nx) {
+                continue;
+            }
+
+            let new_steps = if *d == direction { steps + 1 } else { 1 };
+
+            if new_steps > max_steps {
+                continue;
+            }
+
+            let state = State {
+                heat: heat + grid[ny][nx],
+                position: (ny, nx),
+                direction: *d,
+                steps: new_steps,
+            };
+
+            if *heatmap
+                .get(&((ny, nx), *d, new_steps))
+                .unwrap_or(&usize::MAX)
+                > state.heat
+            {
+                heatmap.insert(((ny, nx), *d, new_steps), state.heat);
+                heap.push(state);
+            }
+        }
+    }
+
+    None
+}
+
+pub fn part_one(input: &str) -> Option<usize> {
+    let grid: Vec<Vec<_>> = input
+        .lines()
+        .map(|x| {
+            x.chars()
+                .filter_map(|x| x.to_digit(10).map(|x| x as usize))
+                .collect()
+        })
+        .collect();
+
+    let target = (grid.len() - 1, grid[0].len() - 1);
+
+    find_path(&grid, (0, 0), target, 0, 3)
+}
+
+pub fn part_two(input: &str) -> Option<usize> {
+    let grid: Vec<Vec<_>> = input
+        .lines()
+        .map(|x| {
+            x.chars()
+                .filter_map(|x| x.to_digit(10).map(|x| x as usize))
+                .collect()
+        })
+        .collect();
+
+    let target = (grid.len() - 1, grid[0].len() - 1);
+
+    find_path(&grid, (0, 0), target, 4, 10)
+}
+
+aoc::solution!(17);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_part_one() {
+        assert_eq!(
+            part_one(&aoc::template::read_file("examples", 17)),
+            Some(102)
+        );
+    }
+    #[test]
+    fn test_part_two() {
+        assert_eq!(
+            part_two(&aoc::template::read_file("examples", 17)),
+            Some(94)
+        );
+    }
+}