solution: day 21

data/examples/21.txt

@@ -0,0 +1,11 @@
+...........
+.....###.#.
+.###.##..#.
+..#.#...#..
+....#.#....
+.##..S####.
+.##..#...#.
+.......##..
+.##.#.####.
+.##..##.##.
+...........

src/bin/21.rs

@@ -0,0 +1,106 @@
+use std::{collections::HashSet, usize};
+
+fn parse_input(input: &str) -> (Vec<Vec<u8>>, (isize, isize)) {
+    let mut grid: Vec<_> = input.lines().map(|x| x.as_bytes().to_vec()).collect();
+
+    let start = grid
+        .iter()
+        .enumerate()
+        .flat_map(|(y, item)| item.iter().enumerate().map(move |(x, item)| (x, y, *item)))
+        .find_map(|(x, y, s)| if s == b'S' { Some((y, x)) } else { None })
+        .unwrap();
+
+    grid[start.0][start.1] = b'.';
+
+    (grid, (start.0 as isize, start.1 as isize))
+}
+
+fn walk_return_at(
+    grid: &[Vec<u8>],
+    start: (isize, isize),
+    mut returns: Vec<usize>,
+    can_cycle: bool,
+) -> Vec<usize> {
+    returns.sort_by(|a, b| b.cmp(a));
+
+    let h = grid.len() as isize;
+    let w = grid[0].len() as isize;
+
+    let invalid_indexing = |y, x| y < 0 || y >= h || x < 0 || x >= w;
+
+    let mut results = Vec::new();
+    let length = returns[0];
+    let mut next = returns.pop().unwrap();
+
+    let mut visited = HashSet::new();
+    visited.insert(start);
+
+    for i in 1..=length {
+        let mut new_visited = HashSet::new();
+
+        for (y, x) in visited.iter() {
+            for (dy, dx) in [(1, 0), (0, 1), (-1, 0), (0, -1)] {
+                let (ny, nx) = (y + dy, x + dx);
+
+                if !can_cycle && invalid_indexing(ny, nx) {
+                    continue;
+                }
+
+                let (cy, cx) = (ny.rem_euclid(h) as usize, nx.rem_euclid(w) as usize);
+
+                if grid[cy][cx] == b'.' {
+                    new_visited.insert((ny, nx));
+                }
+            }
+        }
+
+        visited = new_visited;
+
+        if i == next {
+            results.push(visited.len());
+            if !returns.is_empty() {
+                next = returns.pop().unwrap();
+            }
+        }
+    }
+
+    results
+}
+
+pub fn part_one(input: &str) -> Option<usize> {
+    let (grid, start) = parse_input(input);
+    let result = walk_return_at(&grid, start, vec![64], false);
+
+    Some(result[0])
+}
+
+pub fn part_two(input: &str) -> Option<usize> {
+    let (grid, start) = parse_input(input);
+
+    let h = grid.len();
+    let s = start.0 as usize;
+
+    let result = walk_return_at(&grid, start, vec![s, s + h, s + 2 * h], true);
+
+    let a = (result[2] - 2 * result[1] + result[0]) / 2;
+    let b = (result[1] - result[0]) - a;
+    let c = result[0];
+
+    let x = 26501365 / h;
+
+    Some(a * x * x + b * x + c)
+}
+
+aoc::solution!(21);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn test_part_one() {
+        assert_eq!(
+            part_one(&aoc::template::read_file("examples", 21)),
+            Some(42)
+        );
+    }
+}