Day 16 Complete :) More readable now as well

src/solutions/day16.rs

@@ -2,6 +2,14 @@ use super::Solution;
 
 pub struct Day16 {}
 
+#[derive(Clone, Copy, Debug)]
+enum Direction {
+    Up,
+    Right,
+    Down,
+    Left
+}
+
 impl Solution for Day16 {
     fn part1(
         &self,
@@ -9,14 +17,25 @@ impl Solution for Day16 {
     ) -> Result<Box<dyn std::fmt::Display + Sync>, Box<dyn std::error::Error>> {
         let grid = input.iter().map(|l| l.as_bytes()).collect::<Vec<_>>();
 
-        Ok(Box::new(self.traverse_grid(&grid, (0, 0, 1))))
+        Ok(Box::new(self.traverse_grid(&grid, (0, 0, Direction::Right))))
     }
 
     fn part2(
         &self,
-        _input: &mut Vec<String>,
+        input: &mut Vec<String>,
     ) -> Result<Box<dyn std::fmt::Display + Sync>, Box<dyn std::error::Error>> {
-        Ok(Box::new("Ready"))
+        let grid = input.iter().map(|l| l.as_bytes()).collect::<Vec<_>>();
+
+        // Itterate through every edge piece and record the max energy from a list of all of the combinations
+        // O(n^2) :( Can't think of a better way of processing this
+        // r: row, c: col
+        let ans = (0..grid.len()).flat_map(|r| [(r,0,Direction::Right), (r,grid[0].len()-1,Direction::Left)])
+            .chain((0..grid[0].len()).flat_map(|c| [(0,c,Direction::Down), (grid.len()-1,c,Direction::Up)]))
+            .map(|start| self.traverse_grid(&grid, start))
+            .max()
+            .unwrap();
+
+        Ok(Box::new(ans))
     }
 
     fn get_day(&self) -> u8 {
@@ -26,17 +45,21 @@ impl Solution for Day16 {
 
 impl Day16 {
     /// Returns the next coordinate after taking a step in the given direction
-    fn step(&self, r: usize, c: usize, d: usize) -> (usize, usize, usize) {
-        let (dr, dc) = [(-1,0),(0,1),(1,0),(0,-1)][d];
+    fn step(&self, r: usize, c: usize, d: Direction) -> (usize, usize, Direction) {
+        let (dr, dc) = [(-1,0),(0,1),(1,0),(0,-1)][d as usize];
         ((r as isize + dr) as _, (c as isize + dc) as _, d)
     }
     
     /// Function which traverses through the grid and returns the number of tiles that are energised
     /// in the process
-    fn traverse_grid(&self, grid: &[&[u8]], start: (usize,usize,usize)) -> usize {
-        let mut seen = vec![vec![[false; 4]; grid[0].len()]; grid.len()];
+    fn traverse_grid(&self, grid: &[&[u8]], start: (usize,usize, Direction)) -> usize {
+        let mut energised = vec![vec![[false; 4]; grid[0].len()]; grid.len()];
         let mut beams = vec![start];
 
+        // Direction change given the current direction as an index
+        let right_lean = [Direction::Right, Direction::Up, Direction::Left, Direction::Down];
+        let left_lean = [Direction::Left, Direction::Down, Direction::Right, Direction::Up];
+
         // Trace all of the beams present in the list
         while !beams.is_empty() {
             let mut new_beams = Vec::with_capacity(beams.capacity());
@@ -45,23 +68,25 @@ impl Day16 {
                 if r >= grid.len() || c >= grid[0].len() {
                     continue;
                 }
-                if seen[r][c][d] {
+                if energised[r][c][d as usize] {
                     continue;
                 }
-                seen[r][c][d] = true;
+                energised[r][c][d as usize] = true;
 
                 // Trace the new path in a given direction and draw all of the new beams that come from it.
                 match (grid[r][c], d) {
-                    (b'/',   _) => new_beams.push(self.step(r,c,[1,0,3,2][d])),
-                    (b'\\',  _) => new_beams.push(self.step(r,c,[3,2,1,0][d])),
-                    (b'|', 1|3) => new_beams.extend([self.step(r,c,0), self.step(r,c,2)]),
-                    (b'-', 0|2) => new_beams.extend([self.step(r,c,1), self.step(r,c,3)]),
+                    (b'/',   _) => new_beams.push(self.step(r,c,right_lean[d as usize])),
+                    (b'\\',  _) => new_beams.push(self.step(r,c,left_lean[d as usize])),
+                    (b'|', Direction::Left|Direction::Right) => new_beams.extend([self.step(r,c,Direction::Up), self.step(r,c,Direction::Down)]),
+                    (b'-', Direction::Up|Direction::Down) => new_beams.extend([self.step(r,c,Direction::Left), self.step(r,c, Direction::Right)]),
                     _           => new_beams.push(self.step(r,c,d)),
                 }
             }
             beams = new_beams;
         }
-        seen.iter().flat_map(|row| row).filter(|x| x.iter().any(|&b| b)).count()
+
+        // Get the number of 'energised' tiles now that we have finished processing every beam
+        energised.iter().flat_map(|row| row).filter(|x| x.iter().any(|&b| b)).count()
       }
 }