2021: day04: part 1

aoc2021/src/day04.rs

@@ -0,0 +1,150 @@
+use std::collections::HashMap;
+use std::fmt::Write;
+
+use anyhow::{anyhow, Context, Result};
+
+const INPUT: &str = include_str!("../input/day04.txt");
+
+const GRID_WIDTH: usize = 5;
+const GRID_HEIGHT: usize = 5;
+const GRID_SIZE: usize = GRID_WIDTH * GRID_HEIGHT;
+
+pub fn run() -> Result<String> {
+    let mut res = String::with_capacity(128);
+
+    writeln!(res, "part 1: {}", part1(INPUT)?)?;
+
+    Ok(res)
+}
+
+fn part1(input: &str) -> Result<u64> {
+    let (draws, grids) = input
+        .split_once("\n\n")
+        .context("couldn't split draws from grids")?;
+
+    let draws = draws
+        .split(',')
+        .map(|num| num.parse::<u8>().context("couldn't parse drawn number:"))
+        .collect::<Result<Vec<_>>>()?;
+    let mut grids = grids
+        .split("\n\n")
+        .map(str::parse::<Grid>)
+        .collect::<Result<Vec<_>>>()?;
+
+    let (mut wdraw, mut wgrid) = (None, None);
+
+    'draw_loop: for draw in draws {
+        for grid in &mut grids {
+            if grid.mark(draw) && grid.is_winning() {
+                wgrid = Some(grid.clone());
+                wdraw = Some(draw);
+                break 'draw_loop;
+            }
+        }
+    }
+
+    match (wdraw, wgrid) {
+        (Some(draw), Some(grid)) => {
+            Ok(draw as u64 * grid.unmarked_numbers().map(|n| *n as u64).sum::<u64>())
+        }
+        _ => Err(anyhow!("couldn't find a winning grid!")),
+    }
+}
+
+#[derive(Debug, Clone)]
+struct Grid {
+    number_to_pos: HashMap<u8, (usize, usize)>,
+    pos_to_number: HashMap<(usize, usize), u8>,
+    grid: [bool; GRID_SIZE],
+}
+
+impl Grid {
+    fn mark(&mut self, draw: u8) -> bool {
+        match self.number_to_pos.get(&draw) {
+            Some(&(x, y)) => {
+                *self.access_grid_mut(x, y) = true;
+                true
+            }
+            None => false,
+        }
+    }
+
+    fn is_winning(&self) -> bool {
+        let mut rows = [0u8; GRID_HEIGHT];
+        let mut cols = [0u8; GRID_WIDTH];
+
+        for (y, row) in rows.iter_mut().enumerate() {
+            for (x, col) in cols.iter_mut().enumerate() {
+                if self.access_grid(x, y) {
+                    *row += 1;
+                    *col += 1;
+
+                    if *row as usize == GRID_WIDTH || *col as usize == GRID_HEIGHT {
+                        return true;
+                    }
+                }
+            }
+        }
+
+        false
+    }
+
+    fn unmarked_numbers(&self) -> impl Iterator<Item = &u8> {
+        self.number_to_pos
+            .iter()
+            .filter_map(|(num, &(x, y))| (!self.access_grid(x, y)).then(|| num))
+    }
+
+    fn access_grid(&self, x: usize, y: usize) -> bool {
+        self.grid[y * GRID_HEIGHT + x]
+    }
+
+    fn access_grid_mut(&mut self, x: usize, y: usize) -> &mut bool {
+        &mut self.grid[y * GRID_HEIGHT + x]
+    }
+}
+
+impl std::str::FromStr for Grid {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self> {
+        let mut numbers = s.split_whitespace().map(str::parse);
+        let mut number_to_pos = HashMap::new();
+        let mut pos_to_number = HashMap::new();
+
+        for y in 0..GRID_HEIGHT {
+            for x in 0..GRID_WIDTH {
+                let pos = (x, y);
+                let number: u8 = numbers
+                    .next()
+                    .context("not enough numbers for grid")?
+                    .context("couldn't parse number:")?;
+                number_to_pos.insert(number, pos);
+                pos_to_number.insert(pos, number);
+            }
+        }
+
+        Ok(Grid {
+            number_to_pos,
+            pos_to_number,
+            grid: [false; GRID_SIZE],
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const PROVIDED: &str = include_str!("../input/day04_provided.txt");
+
+    #[test]
+    fn part1_provided() {
+        assert_eq!(part1(PROVIDED).unwrap(), 4512);
+    }
+
+    #[test]
+    fn part1_real() {
+        assert_eq!(part1(INPUT).unwrap(), 45031);
+    }
+}

aoc2021/src/lib.rs

@@ -3,3 +3,4 @@
 pub mod day01;
 pub mod day02;
 pub mod day03;
+pub mod day04;

aoc2021/src/main.rs

@@ -5,9 +5,10 @@ use aoc::DayFunc;
 use aoc2021::day01;
 use aoc2021::day02;
 use aoc2021::day03;
+use aoc2021::day04;
 
 fn main() -> Result<()> {
-    let days: &[DayFunc] = &[day01::run, day02::run, day03::run];
+    let days: &[DayFunc] = &[day01::run, day02::run, day03::run, day04::run];
 
     aoc::run(days)
 }