Refactor day 4 to use const generics

2021/src/bin/day04.rs

@@ -1,31 +1,52 @@
 #![feature(test)]
 extern crate test;
+use std::hash::Hash;
+
 use aoc2021::common::*;
 use itertools::Itertools;
 
 const DAY: usize = 4;
 const BOARD_SIZE: usize = 5;
 const NUMBERS_PER_BOARD: usize = BOARD_SIZE * BOARD_SIZE;
-type Board = Vec<Vec<u8>>;
+type Board<T> = [[T; BOARD_SIZE]; BOARD_SIZE * 2];
+
+const WINNING_INDICES: Board<usize> = {
+    let mut out = [[0; BOARD_SIZE]; BOARD_SIZE * 2];
+    let mut i = 0;
+    while i < BOARD_SIZE {
+        // tfw there are no for loops in const fn
+        let mut j = 0;
+        while j < BOARD_SIZE {
+            out[i][j] = i * 5 + j;
+            out[BOARD_SIZE + i][j] = i + j * 5;
+            j += 1;
+        }
+        i += 1;
+    }
+    out
+};
 
 #[derive(Debug, Clone)]
-struct BingoGame {
-    input_numbers: Vec<u8>,
-    boards:        Vec<Board>,
+struct BingoGame<T> {
+    input_numbers: Vec<T>,
+    boards:        Vec<Board<Option<T>>>,
 }
 
-impl BingoGame {
-    fn mark_number(&mut self, n: &u8) {
+impl<T: Eq> BingoGame<T> {
+    fn mark_number(&mut self, n: &T) {
         for board in self.boards.iter_mut() {
             // Actually using sets instead of vectors here takes twice as long,
             // so just pretend these are actually sets.
             for winning_set in board.iter_mut() {
-                winning_set.retain(|e| e != n);
+                // winning_set.retain(|e| e != n);
+                for w in winning_set.iter_mut().filter(|w| w.as_ref() == Some(n)) {
+                    *w = None
+                }
             }
         }
     }
 
-    fn find_winner(&self) -> Option<&Board> {
+    fn find_winner(&self) -> Option<&Board<Option<T>>> {
         self.boards.iter().find(|b| has_won(b))
     }
 
@@ -36,11 +57,11 @@ impl BingoGame {
     }
 }
 
-fn has_won(board: &Board) -> bool {
-    board.iter().any(|s| s.is_empty())
+fn has_won<T>(board: &Board<Option<T>>) -> bool {
+    board.iter().any(|s| s.iter().all(|w| w.is_none()))
 }
 
-fn parse_input(raw: &str) -> BingoGame {
+fn parse_input(raw: &str) -> BingoGame<u8> {
     let (input_numbers, boards) = raw.split_once("\n\n").unwrap();
     let input_numbers = input_numbers.split(',').map(|n| n.parse().unwrap()).collect();
     let boards = boards
@@ -48,25 +69,18 @@ fn parse_input(raw: &str) -> BingoGame {
         .map(|b| b.split_ascii_whitespace().map(|n| n.parse().unwrap()).collect())
         .map(|v: Vec<u8>| {
             debug_assert_eq!(v.len(), NUMBERS_PER_BOARD);
-            // This seems way too complicated. What am I missing?
-            (0..NUMBERS_PER_BOARD)
-                .chain((0..BOARD_SIZE).flat_map(|i| (i..NUMBERS_PER_BOARD).step_by(BOARD_SIZE)))
-                .map(|i| v[i])
-                .chunks(BOARD_SIZE)
-                .into_iter()
-                .map(|c| c.collect())
-                .collect()
+            WINNING_INDICES.map(|row_or_col| row_or_col.map(|i| Some(v[i])))
         })
         .collect();
     BingoGame { input_numbers, boards }
 }
 
-fn board_score(board: &Board, current_number: u8) -> usize {
-    let remainder: usize = board.iter().flatten().unique().map(|&n| n as usize).sum();
+fn board_score<T: Hash + Eq + Into<usize> + Copy>(board: &Board<Option<T>>, current_number: u8) -> usize {
+    let remainder: usize = board.iter().flatten().flatten().unique().map::<usize, _>(|&n| n.into()).sum();
     remainder * (current_number as usize)
 }
 
-fn part1(parsed: &BingoGame) -> usize {
+fn part1(parsed: &BingoGame<u8>) -> usize {
     let mut game = parsed.to_owned();
     for n in &game.input_numbers.clone() {
         game.mark_number(n);
@@ -77,7 +91,7 @@ fn part1(parsed: &BingoGame) -> usize {
     unreachable!("Game should have ended at some point")
 }
 
-fn part2(parsed: &BingoGame) -> usize {
+fn part2(parsed: &BingoGame<u8>) -> usize {
     let mut game = parsed.to_owned();
     for n in &game.input_numbers.clone() {
         game.mark_number(n);
@@ -91,7 +105,7 @@ fn part2(parsed: &BingoGame) -> usize {
 
 fn main() {
     let input = parse_input(&read_file(DAY));
-    println!("{input:?}");
+    println!("{WINNING_INDICES:?}");
     println!("Part 1: {}", part1(&input));
     println!("Part 2: {}", part2(&input));
 }