1  #![feature(test)]
  2  
  3  use rustc_hash::FxHashSet;
  4  
  5  struct Input {
  6      nums: Vec<i32>,
  7      boards: Vec<Vec<Vec<i32>>>,
  8  }
  9  
 10  fn setup(input: &str) -> Input {
 11      let nums = input
 12          .lines()
 13          .next()
 14          .unwrap()
 15          .split(',')
 16          .map(|n| n.parse().unwrap())
 17          .collect();
 18      let boards = input
 19          .split("\n\n")
 20          .skip(1)
 21          .map(|board| {
 22              board
 23                  .lines()
 24                  .map(|line| {
 25                      line.split_whitespace()
 26                          .map(|n| n.parse().unwrap())
 27                          .collect()
 28                  })
 29                  .collect()
 30          })
 31          .collect();
 32      Input { nums, boards }
 33  }
 34  
 35  struct State {
 36      boards: Vec<Vec<Vec<i32>>>,
 37      marked: FxHashSet<i32>,
 38  }
 39  
 40  impl State {
 41      fn from_input(input: &Input) -> State {
 42          State {
 43              boards: input.boards.clone(),
 44              marked: FxHashSet::default(),
 45          }
 46      }
 47  
 48      fn mark(&mut self, num: i32) {
 49          self.marked.insert(num);
 50      }
 51  
 52      fn check_board(&self, board: usize) -> bool {
 53          let board = &self.boards[board];
 54          for row in board {
 55              if row.iter().all(|n| self.marked.contains(n)) {
 56                  return true;
 57              }
 58          }
 59          for i in 0..board[0].len() {
 60              if board.iter().all(|row| self.marked.contains(&row[i])) {
 61                  return true;
 62              }
 63          }
 64          false
 65      }
 66  
 67      fn get_score(&self, board: usize) -> i32 {
 68          self.boards[board]
 69              .iter()
 70              .flatten()
 71              .filter(|n| !self.marked.contains(n))
 72              .copied()
 73              .reduce(|a, b| a + b)
 74              .unwrap()
 75      }
 76  }
 77  
 78  fn part1(input: &Input) -> String {
 79      let mut state = State::from_input(input);
 80      for num in &input.nums {
 81          state.mark(*num);
 82          for i in 0..state.boards.len() {
 83              if state.check_board(i) {
 84                  return (num * state.get_score(i)).to_string();
 85              }
 86          }
 87      }
 88      panic!();
 89  }
 90  
 91  fn part2(input: &Input) -> String {
 92      let mut state = State::from_input(input);
 93      for num in &input.nums {
 94          state.mark(*num);
 95          let mut i = 0;
 96          while i < state.boards.len() {
 97              if !state.check_board(i) {
 98                  i += 1;
 99                  continue;
100              }
101              if state.boards.len() == 1 {
102                  return (num * state.get_score(i)).to_string();
103              }
104              state.boards.remove(i);
105          }
106      }
107      panic!()
108  }
109  
110  aoc::main!(2021, 4, ex: 1);