06.rs
1 #![feature(test)] 2 3 use aoc::{grid::Direction, iter_ext::IterExt}; 4 use rayon::iter::{IntoParallelIterator, ParallelIterator}; 5 use rustc_hash::FxHashSet; 6 7 type Input = Vec<Vec<Cell>>; 8 9 #[derive(Debug, Clone, Copy, PartialEq, Eq)] 10 enum Cell { 11 Guard, // ^ 12 Obstruction, // # 13 Empty, // . 14 } 15 16 fn setup(input: &str) -> Input { 17 input 18 .trim() 19 .lines() 20 .map(|line| { 21 line.trim() 22 .bytes() 23 .map(|b| match b { 24 b'^' => Cell::Guard, 25 b'#' => Cell::Obstruction, 26 b'.' => Cell::Empty, 27 _ => panic!(), 28 }) 29 .collect() 30 }) 31 .collect() 32 } 33 34 #[derive(Debug)] 35 struct Walk<'a> { 36 x: usize, 37 y: usize, 38 d: Direction, 39 grid: &'a Input, 40 init: bool, 41 } 42 43 impl Walk<'_> { 44 fn from_input(input: &Input) -> Walk { 45 let (x, y) = input 46 .iter() 47 .enumerate() 48 .find_map(|(y, row)| row.iter().index(&Cell::Guard).map(|x| (x, y))) 49 .unwrap(); 50 Walk { 51 x, 52 y, 53 d: Direction::North, 54 grid: input, 55 init: false, 56 } 57 } 58 } 59 60 impl Iterator for Walk<'_> { 61 type Item = (usize, usize, Direction); 62 63 fn next(&mut self) -> Option<Self::Item> { 64 if !self.init { 65 self.init = true; 66 return Some((self.x, self.y, self.d)); 67 } 68 69 loop { 70 let (nx, ny) = self 71 .d 72 .step(self.x, self.y, self.grid[0].len(), self.grid.len())?; 73 if !matches!(self.grid[ny][nx], Cell::Obstruction) { 74 (self.x, self.y) = (nx, ny); 75 break; 76 } 77 self.d = self.d.rotate_right(); 78 } 79 80 Some((self.x, self.y, self.d)) 81 } 82 } 83 84 #[derive(Debug)] 85 struct Grid { 86 width: usize, 87 height: usize, 88 start: (usize, usize), 89 north: Vec<LookupRow>, 90 east: Vec<LookupRow>, 91 south: Vec<LookupRow>, 92 west: Vec<LookupRow>, 93 } 94 95 #[derive(Debug)] 96 struct LookupRow(Vec<usize>); 97 98 impl LookupRow { 99 fn lookup(&self, n: usize, extra: Option<usize>) -> usize { 100 match extra { 101 Some(x) if x <= n => self.0[n].min(n - x), 102 _ => self.0[n], 103 } 104 } 105 } 106 107 impl FromIterator<bool> for LookupRow { 108 fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self { 109 Self( 110 iter.into_iter() 111 .scan(usize::MAX, |s, x| { 112 *s = if x { 0 } else { s.saturating_add(1) }; 113 Some(*s) 114 }) 115 .collect(), 116 ) 117 } 118 } 119 120 impl Grid { 121 fn from_input(input: &Input) -> Self { 122 let width = input[0].len(); 123 let height = input.len(); 124 125 let start = input 126 .iter() 127 .enumerate() 128 .find_map(|(y, row)| row.iter().index(&Cell::Guard).map(|x| (x, y))) 129 .unwrap(); 130 131 let is_obstruction = |x: usize, y: usize| matches!(input[y][x], Cell::Obstruction); 132 133 Self { 134 width, 135 height, 136 start, 137 north: (0..width) 138 .into_par_iter() 139 .map(|x| (0..height).map(|y| is_obstruction(x, y)).collect()) 140 .collect(), 141 east: (0..height) 142 .into_par_iter() 143 .map(|y| (0..width).rev().map(|x| is_obstruction(x, y)).collect()) 144 .collect(), 145 south: (0..width) 146 .into_par_iter() 147 .map(|x| (0..height).rev().map(|y| is_obstruction(x, y)).collect()) 148 .collect(), 149 west: (0..height) 150 .into_par_iter() 151 .map(|y| (0..width).map(|x| is_obstruction(x, y)).collect()) 152 .collect(), 153 } 154 } 155 156 fn check_cycle(&self, extra_obstruction: (usize, usize)) -> bool { 157 let (mut x, mut y) = self.start; 158 let mut d = Direction::North; 159 let mut visited = FxHashSet::default(); 160 while visited.insert((x, y, d)) { 161 let Some((nx, ny)) = self.walk(x, y, d, extra_obstruction) else { 162 return false; 163 }; 164 (x, y, d) = (nx, ny, d.rotate_right()) 165 } 166 true 167 } 168 169 fn walk( 170 &self, 171 x: usize, 172 y: usize, 173 direction: Direction, 174 extra_obstruction: (usize, usize), 175 ) -> Option<(usize, usize)> { 176 let eo = Some(extra_obstruction); 177 Some(match direction { 178 Direction::North => self.north[x].lookup(y, eo.filter(|eo| eo.0 == x).map(|eo| eo.1)), 179 Direction::East => self.east[y].lookup( 180 self.width - x - 1, 181 eo.filter(|eo| eo.1 == y).map(|eo| self.width - eo.0 - 1), 182 ), 183 Direction::South => self.south[x].lookup( 184 self.height - y - 1, 185 eo.filter(|eo| eo.0 == x).map(|eo| self.height - eo.1 - 1), 186 ), 187 Direction::West => self.west[y].lookup(x, eo.filter(|eo| eo.1 == y).map(|eo| eo.0)), 188 }) 189 .filter(|&r| r != usize::MAX) 190 .map(|r| { 191 let r = r - 1; 192 match direction { 193 Direction::North => (x, y - r), 194 Direction::East => (x + r, y), 195 Direction::South => (x, y + r), 196 Direction::West => (x - r, y), 197 } 198 }) 199 } 200 } 201 202 fn part1(input: &Input) -> usize { 203 Walk::from_input(input) 204 .map(|(x, y, _)| (x, y)) 205 .collect::<FxHashSet<_>>() 206 .len() 207 } 208 209 fn part2(input: &Input) -> usize { 210 let candidates = Walk::from_input(input) 211 .map(|(x, y, _)| (x, y)) 212 .filter(|&(x, y)| matches!(input[y][x], Cell::Empty)) 213 .collect::<FxHashSet<_>>(); 214 215 let grid = Grid::from_input(input); 216 217 candidates 218 .into_par_iter() 219 .filter(|&(ox, oy)| grid.check_cycle((ox, oy))) 220 .count() 221 } 222 223 aoc::main!(2024, 6, ex: 1);