2020: day11: part 2

aoc2020/src/day11.rs

@@ -1,5 +1,4 @@
 use std::fmt::Write;
-use std::mem;
 
 use aoc::err;
 
@@ -9,6 +8,7 @@ pub fn run() -> aoc::Result<String> {
     let mut res = String::with_capacity(128);
 
     writeln!(res, "part 1: {}", part1(INPUT)?)?;
+    writeln!(res, "part 2: {}", part2(INPUT)?)?;
 
     Ok(res)
 }
@@ -16,7 +16,17 @@ pub fn run() -> aoc::Result<String> {
 fn part1(input: &str) -> aoc::Result<usize> {
     let mut layout: Layout = input.parse()?;
 
-    layout.converge();
+    let occupied_threshold = 4;
+    layout.converge(occupied_threshold, Layout::count_adjacent);
+
+    Ok(layout.occupied_seats())
+}
+
+fn part2(input: &str) -> aoc::Result<usize> {
+    let mut layout: Layout = input.parse()?;
+
+    let occupied_threshold = 5;
+    layout.converge(occupied_threshold, Layout::count_line_of_sight);
 
     Ok(layout.occupied_seats())
 }
@@ -38,10 +48,15 @@ struct Layout {
 
 impl Layout {
     /// Steps one round in the simulation, returns the previous grid
-    fn step(&mut self) -> Grid {
+    fn step(
+        &mut self,
+        occupied_threshold: u8,
+        adj_count: fn(&Self, usize, usize, Cell) -> u8,
+    ) -> bool {
         let grid = &self.grid;
 
         let mut new = grid.clone();
+        let mut changed = false;
 
         for i in 0..self.height {
             for j in 0..self.width {
@@ -49,13 +64,15 @@ impl Layout {
 
                 match cell {
                     Cell::EmptySeat => {
-                        if self.count_adjacent(i, j, Cell::OccupiedSeat) == 0 {
+                        if adj_count(&self, i, j, Cell::OccupiedSeat) == 0 {
                             new[i][j] = Cell::OccupiedSeat;
+                            changed = true;
                         }
                     }
                     Cell::OccupiedSeat => {
-                        if self.count_adjacent(i, j, Cell::OccupiedSeat) >= 4 {
+                        if adj_count(&self, i, j, Cell::OccupiedSeat) >= occupied_threshold {
                             new[i][j] = Cell::EmptySeat;
+                            changed = true;
                         }
                     }
                     _ => {}
@@ -63,16 +80,14 @@ impl Layout {
             }
         }
 
-        mem::replace(&mut self.grid, new)
+        self.grid = new;
+
+        changed
     }
 
     /// Steps through the simulation until a fixpoint is reached
-    fn converge(&mut self) {
-        let mut prev = self.step();
-
-        while prev != self.grid {
-            prev = self.step();
-        }
+    fn converge(&mut self, occupied_threshold: u8, adj_count: fn(&Self, usize, usize, Cell) -> u8) {
+        while self.step(occupied_threshold, adj_count) {}
     }
 
     const OFFSETS: &'static [(i8, i8)] = &[
@@ -105,6 +120,36 @@ impl Layout {
         count
     }
 
+    fn count_line_of_sight(&self, i: usize, j: usize, value: Cell) -> u8 {
+        let mut count = 0;
+
+        for (di, dj) in Self::OFFSETS {
+            let mut distance = 1;
+
+            let diff = loop {
+                let (di, dj) = (di * distance, dj * distance);
+
+                let (i, j) = (i.wrapping_add(di as usize), j.wrapping_add(dj as usize));
+
+                let cell = self.grid.get(i).map(|line| line.get(j)).flatten();
+
+                match cell {
+                    // keep going, the next seat is farther away
+                    Some(Cell::Floor) => distance += 1,
+                    // found the kind of seat we care about
+                    Some(&seat) if seat == value => break 1,
+                    // found a seat that blocks line of sight, or reached out of bounds
+                    Some(_) | None => break 0,
+                }
+            };
+
+            // only check seat if it's in bounds
+            count += diff;
+        }
+
+        count
+    }
+
     fn occupied_seats(&self) -> usize {
         self.grid
             .iter()
@@ -175,4 +220,14 @@ mod tests {
     fn part1_real() {
         assert_eq!(part1(INPUT).unwrap(), 2427);
     }
+
+    #[test]
+    fn part2_provided() {
+        assert_eq!(part2(PROVIDED).unwrap(), 26);
+    }
+
+    #[test]
+    fn part2_real() {
+        assert_eq!(part2(INPUT).unwrap(), 2199);
+    }
 }