1  #![feature(test)]
 2  
 3  use itertools::{FoldWhile::Continue, Itertools};
 4  use rustc_hash::FxHashMap;
 5  
 6  type Input = Vec<Vec<Cell>>;
 7  
 8  #[derive(Debug, Clone, Copy)]
 9  enum Cell {
10      Digit(u8),
11      Symbol { gear: bool },
12      Empty,
13  }
14  
15  fn setup(input: &str) -> Input {
16      input
17          .lines()
18          .map(|line| {
19              line.bytes()
20                  .map(|b| match b {
21                      b'0'..=b'9' => Cell::Digit(b - b'0'),
22                      b'.' => Cell::Empty,
23                      b'*' => Cell::Symbol { gear: true },
24                      _ => Cell::Symbol { gear: false },
25                  })
26                  .collect()
27          })
28          .collect()
29  }
30  
31  fn find_part_numbers(
32      input: &Input,
33  ) -> impl Iterator<Item = (u32, impl Iterator<Item = (usize, usize)> + '_)> + '_ {
34      input.iter().enumerate().flat_map(move |(i, line)| {
35          line.iter()
36              .copied()
37              .enumerate()
38              .batching(|it| {
39                  let (num, j) = it.find_map(|(j, c)| match c {
40                      Cell::Digit(d) => Some((d as u32, j)),
41                      _ => None,
42                  })?;
43                  let (num, k) = it
44                      .fold_while((num, j), |acc @ (num, _), (k, c)| {
45                          match match c {
46                              Cell::Digit(d) => Some(d),
47                              _ => None,
48                          } {
49                              Some(d) => Continue((num * 10 + d as u32, k)),
50                              None => itertools::FoldWhile::Done(acc),
51                          }
52                      })
53                      .into_inner();
54                  Some((num, j, k))
55              })
56              .map(move |(num, j, k)| {
57                  let gears = (i.saturating_sub(1)..=(i + 1).min(input.len() - 1))
58                      .dedup()
59                      .flat_map(move |i| {
60                          (j.saturating_sub(1)..=(k + 1).min(input[i].len() - 1))
61                              .filter(move |&j| matches!(input[i][j], Cell::Symbol { .. }))
62                              .map(move |j| (i, j))
63                      });
64                  (num, gears)
65              })
66      })
67  }
68  
69  fn part1(input: &Input) -> u32 {
70      find_part_numbers(input)
71          .filter_map(|(num, mut parts)| parts.next().map(|_| num))
72          .sum()
73  }
74  
75  fn part2(input: &Input) -> u32 {
76      let mut gears = FxHashMap::<(usize, usize), _>::default();
77      for (num, gs) in find_part_numbers(input) {
78          for g in gs {
79              gears.entry(g).or_insert_with(Vec::new).push(num);
80          }
81      }
82      gears
83          .into_iter()
84          .filter(|&((i, j), ref nums)| {
85              matches!(input[i][j], Cell::Symbol { gear: true }) && nums.len() == 2
86          })
87          .map(|(_, nums)| nums.into_iter().product::<u32>())
88          .sum()
89  }
90  
91  aoc::main!(2023, 3, ex: 1);