From b51e5b8ef272a5eb9fc19dc107480282ab2df285 Mon Sep 17 00:00:00 2001
From: Connor Slade <connor@connorcode.com>
Date: Mon, 16 Dec 2024 16:08:03 -0500
Subject: [PATCH] [2024] Cleanup day 16

---
 aoc_2024/src/day_16.rs            | 285 +++++++++++-------------------
 aoc_lib/src/direction/cardinal.rs |   1 +
 aoc_lib/src/matrix.rs             |  19 +-
 3 files changed, 126 insertions(+), 179 deletions(-)

diff --git a/aoc_2024/src/day_16.rs b/aoc_2024/src/day_16.rs
index c82e689..b650afb 100644
--- a/aoc_2024/src/day_16.rs
+++ b/aoc_2024/src/day_16.rs
@@ -1,11 +1,33 @@
-use std::collections::{BinaryHeap, HashMap, HashSet};
+use std::{
+    cmp::Ordering,
+    collections::{BinaryHeap, HashSet},
+};
 
 use aoc_lib::{direction::cardinal::Direction, matrix::Matrix};
 use common::{solution, Answer};
-use nd_vec::{vector, Vec2};
+use nd_vec::Vec2;
 
 solution!("Reindeer Maze", 16);
 
+fn part_a(input: &str) -> Answer {
+    let map = Maze::parse(input);
+    let (_, shortest) = map.foreword();
+    shortest.into()
+}
+
+fn part_b(input: &str) -> Answer {
+    let map = Maze::parse(input);
+    let (scores, _) = map.foreword();
+    map.reverse(scores).into()
+}
+
+struct Maze {
+    map: Matrix<Tile>,
+
+    start: Vec2<usize>,
+    end: Vec2<usize>,
+}
+
 #[derive(PartialEq, Eq)]
 enum Tile {
     Empty,
@@ -14,210 +36,117 @@ enum Tile {
     End,
 }
 
-fn part_a(input: &str) -> Answer {
-    let map = Matrix::new_chars(input, |c| match c {
-        '.' => Tile::Empty,
-        '#' => Tile::Wall,
-        'S' => Tile::Start,
-        'E' => Tile::End,
-        _ => panic!(),
-    });
-    let start = map.find(Tile::Start).unwrap();
-
-    // find minimum score needed to path from S to E
-    // - Move foreword (+1)
-    // - rotate left or right (+1000)
-
-    #[derive(PartialEq, Eq)]
-    struct Item {
-        pos: Vec2<usize>,
-        dir: Direction,
-        score: u32,
-    }
-
-    impl Ord for Item {
-        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-            other.score.cmp(&self.score)
-        }
-    }
-
-    impl PartialOrd for Item {
-        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-            Some(other.score.cmp(&self.score))
-        }
-    }
+#[derive(PartialEq, Eq)]
+struct Item {
+    pos: Vec2<usize>,
+    dir: Direction,
+    score: u32,
+}
 
-    // todo: try bin heap
-    let mut queue = BinaryHeap::new();
-    let mut seen = HashSet::new();
+impl Maze {
+    fn parse(input: &str) -> Self {
+        let map = Matrix::new_chars(input, |c| match c {
+            '.' => Tile::Empty,
+            '#' => Tile::Wall,
+            'S' => Tile::Start,
+            'E' => Tile::End,
+            _ => panic!(),
+        });
 
-    queue.push(Item {
-        pos: start,
-        dir: Direction::Right,
-        score: 0,
-    });
+        let start = map.find(Tile::Start).unwrap();
+        let end = map.find(Tile::End).unwrap();
 
-    while let Some(Item { pos, dir, score }) = queue.pop() {
-        if !seen.insert((pos, dir)) {
-            continue;
-        }
+        Self { map, start, end }
+    }
 
-        if map[pos] == Tile::End {
-            return score.into();
-        }
+    /// Use dijkstra's to find the shortest path, populating a costs grid with
+    /// the minimum cost needed to reach that tile, which will be used for part B.
+    fn foreword(&self) -> (Matrix<[u32; 4]>, u32) {
+        let mut queue = BinaryHeap::new();
+        let mut seen = HashSet::new();
+        let mut costs = Matrix::new_default(self.map.size, [u32::MAX; 4]);
 
-        // Move foreword
-        let next = dir.wrapping_advance(pos);
-        if map.contains(next) && map[next] != Tile::Wall {
-            queue.push(Item {
-                pos: next,
-                dir,
-                score: score + 1,
-            });
-        }
+        queue.push(Item::new(self.start, Direction::Right, 0));
 
-        queue.push(Item {
-            pos,
-            dir: dir.turn_left(),
-            score: score + 1000,
-        });
+        while let Some(Item { pos, dir, score }) = queue.pop() {
+            let min = &mut costs[pos];
+            min[dir as usize] = min[dir as usize].min(score);
 
-        queue.push(Item {
-            pos,
-            dir: dir.turn_right(),
-            score: score + 1000,
-        });
-    }
+            if !seen.insert((pos, dir)) {
+                continue;
+            }
 
-    unreachable!()
-}
+            if self.map[pos] == Tile::End {
+                return (costs, score);
+            }
 
-fn part_b(input: &str) -> Answer {
-    let map = Matrix::new_chars(input, |c| match c {
-        '.' => Tile::Empty,
-        '#' => Tile::Wall,
-        'S' => Tile::Start,
-        'E' => Tile::End,
-        _ => panic!(),
-    });
-    let start = map.find(Tile::Start).unwrap();
-
-    // find minimum score needed to path from S to E
-    // - Move foreword (+1)
-    // - rotate left or right (+1000)
-
-    #[derive(PartialEq, Eq)]
-    struct Item {
-        pos: Vec2<usize>,
-        dir: Direction,
-        score: u32,
-
-        path: Vec<Vec2<usize>>,
-    }
+            let next = dir.wrapping_advance(pos);
+            if self.map.contains(next) && self.map[next] != Tile::Wall {
+                queue.push(Item::new(next, dir, score + 1));
+            }
 
-    impl Ord for Item {
-        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-            other.score.cmp(&self.score)
+            for dir in [dir.turn_left(), dir.turn_right()] {
+                queue.push(Item::new(pos, dir, score + 1000));
+            }
         }
-    }
 
-    impl PartialOrd for Item {
-        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-            Some(other.score.cmp(&self.score))
-        }
+        unreachable!("No path found")
     }
 
-    // todo: try bin heap
-    let mut queue = BinaryHeap::new();
-    let mut seen = HashMap::new();
-
-    let mut on_best = HashSet::<Vec2<usize>>::new();
+    /// Walks backwards from the end using a BFS to find all the tiles that are
+    /// on any of the shortest path.
+    fn reverse(&self, mut scores: Matrix<[u32; 4]>) -> u32 {
+        let mut seen = HashSet::new();
+        let mut seen_queue = vec![];
 
-    queue.push(Item {
-        pos: start,
-        dir: Direction::Right,
-        score: 0,
-
-        path: vec![start],
-    });
-
-    let mut best = None;
+        let end_lowest = scores.get(self.end).unwrap();
+        for dir in Direction::ALL {
+            let min_cost = end_lowest[dir as usize];
+            seen_queue.push(Item::new(self.end, dir, min_cost));
+        }
 
-    while let Some(Item {
-        pos,
-        dir,
-        score,
-        path,
-    }) = queue.pop()
-    {
-        if let Some(&prev) = seen.get(&(pos, dir)) {
-            if score > prev {
+        while let Some(item) = seen_queue.pop() {
+            seen.insert(item.pos);
+            if item.pos == self.start {
                 continue;
             }
-        } else {
-            seen.insert((pos, dir), score);
-        }
 
-        if map[pos] == Tile::End {
-            if let Some(real_best) = best {
-                dbg!(score, real_best);
-                if score == real_best {
-                    on_best.extend(path.iter());
+            let next = item.dir.opposite().wrapping_advance(item.pos);
+            for next in [
+                Item::new(next, item.dir, item.score - 1),
+                Item::new(item.pos, item.dir.turn_left(), item.score - 1000),
+                Item::new(item.pos, item.dir.turn_right(), item.score - 1000),
+            ] {
+                if self.map.contains(next.pos)
+                    && self.map[next.pos] != Tile::Wall
+                    && next.score == scores.get(next.pos).unwrap()[next.dir as usize]
+                {
+                    scores.get_mut(next.pos).unwrap()[next.dir as usize] = u32::MAX;
+                    seen_queue.push(next);
                 }
-            } else {
-                best = Some(score);
-                on_best.extend(path.iter());
             }
-
-            continue;
-        }
-
-        // Move foreword
-        let next = dir.wrapping_advance(pos);
-        let mut next_path = path.clone();
-        next_path.push(next);
-        if map.contains(next) && map[next] != Tile::Wall {
-            queue.push(Item {
-                pos: next,
-                dir,
-                score: score + 1,
-                path: next_path,
-            });
         }
 
-        queue.push(Item {
-            pos,
-            dir: dir.turn_left(),
-            score: score + 1000,
-            path: path.clone(),
-        });
+        seen.len() as u32
+    }
+}
 
-        queue.push(Item {
-            pos,
-            dir: dir.turn_right(),
-            score: score + 1000,
-            path,
-        });
+impl Item {
+    fn new(pos: Vec2<usize>, dir: Direction, score: u32) -> Self {
+        Self { pos, dir, score }
     }
+}
 
-    for y in 0..map.size.y() {
-        for x in 0..map.size.x() {
-            let pos = vector!(x, y);
-            print!(
-                "{}",
-                match map[pos] {
-                    Tile::Empty if on_best.contains(&pos) => '@',
-                    Tile::Empty => '.',
-                    Tile::Wall => '#',
-                    Tile::Start => 'S',
-                    Tile::End => 'E',
-                }
-            );
-        }
-        println!();
+impl Ord for Item {
+    fn cmp(&self, other: &Self) -> Ordering {
+        other.score.cmp(&self.score)
     }
+}
 
-    on_best.len().into()
+impl PartialOrd for Item {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
 }
 
 #[cfg(test)]
diff --git a/aoc_lib/src/direction/cardinal.rs b/aoc_lib/src/direction/cardinal.rs
index 05180fb..5cb0ddd 100644
--- a/aoc_lib/src/direction/cardinal.rs
+++ b/aoc_lib/src/direction/cardinal.rs
@@ -2,6 +2,7 @@ use nd_vec::{vector, Vec2};
 use num_traits::{Num, Signed, WrappingSub};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[repr(u8)]
 pub enum Direction {
     Up,
     Down,
diff --git a/aoc_lib/src/matrix.rs b/aoc_lib/src/matrix.rs
index 3bf80ad..b2ed519 100644
--- a/aoc_lib/src/matrix.rs
+++ b/aoc_lib/src/matrix.rs
@@ -1,4 +1,8 @@
-use std::{fmt::Debug, hash::Hash, ops::Index};
+use std::{
+    fmt::Debug,
+    hash::Hash,
+    ops::{Index, IndexMut},
+};
 
 use nd_vec::{vector, Vec2};
 use num_traits::{Num, ToPrimitive};
@@ -89,6 +93,12 @@ impl<T> Index<[usize; 2]> for Matrix<T> {
     }
 }
 
+impl<T> IndexMut<[usize; 2]> for Matrix<T> {
+    fn index_mut(&mut self, index: [usize; 2]) -> &mut T {
+        &mut self.data[index[1] * self.size.x() + index[0]]
+    }
+}
+
 impl<T, K: ToPrimitive + Copy> Index<Vec2<K>> for Matrix<T> {
     type Output = T;
 
@@ -98,6 +108,13 @@ impl<T, K: ToPrimitive + Copy> Index<Vec2<K>> for Matrix<T> {
     }
 }
 
+impl<T, K: ToPrimitive + Copy> IndexMut<Vec2<K>> for Matrix<T> {
+    fn index_mut(&mut self, index: Vec2<K>) -> &mut T {
+        let index = index.num_cast::<usize>().unwrap();
+        &mut self.data[index.y() * self.size.x() + index.x()]
+    }
+}
+
 impl<T: Clone> Clone for Matrix<T> {
     fn clone(&self) -> Self {
         Self {