91 lines
2.7 KiB
Rust
91 lines
2.7 KiB
Rust
#![feature(test)]
|
|
|
|
use itertools::{FoldWhile::Continue, Itertools};
|
|
use rustc_hash::FxHashMap;
|
|
|
|
type Input = Vec<Vec<Cell>>;
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
enum Cell {
|
|
Digit(u8),
|
|
Symbol { gear: bool },
|
|
Empty,
|
|
}
|
|
|
|
fn setup(input: &str) -> Input {
|
|
input
|
|
.lines()
|
|
.map(|line| {
|
|
line.bytes()
|
|
.map(|b| match b {
|
|
b'0'..=b'9' => Cell::Digit(b - b'0'),
|
|
b'.' => Cell::Empty,
|
|
b'*' => Cell::Symbol { gear: true },
|
|
_ => Cell::Symbol { gear: false },
|
|
})
|
|
.collect()
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn find_part_numbers(
|
|
input: &Input,
|
|
) -> impl Iterator<Item = (u32, impl Iterator<Item = (usize, usize)> + '_)> + '_ {
|
|
input.iter().enumerate().flat_map(move |(i, line)| {
|
|
line.iter()
|
|
.copied()
|
|
.enumerate()
|
|
.batching(|it| {
|
|
let (num, j) = it.find_map(|(j, c)| match c {
|
|
Cell::Digit(d) => Some((d as u32, j)),
|
|
_ => None,
|
|
})?;
|
|
let (num, k) = it
|
|
.fold_while((num, j), |acc @ (num, _), (k, c)| {
|
|
match match c {
|
|
Cell::Digit(d) => Some(d),
|
|
_ => None,
|
|
} {
|
|
Some(d) => Continue((num * 10 + d as u32, k)),
|
|
None => itertools::FoldWhile::Done(acc),
|
|
}
|
|
})
|
|
.into_inner();
|
|
Some((num, j, k))
|
|
})
|
|
.map(move |(num, j, k)| {
|
|
let gears = (i.saturating_sub(1)..=(i + 1).min(input.len() - 1))
|
|
.dedup()
|
|
.flat_map(move |i| {
|
|
(j.saturating_sub(1)..=(k + 1).min(input[i].len() - 1))
|
|
.filter(move |&j| matches!(input[i][j], Cell::Symbol { .. }))
|
|
.map(move |j| (i, j))
|
|
});
|
|
(num, gears)
|
|
})
|
|
})
|
|
}
|
|
|
|
fn part1(input: &Input) -> u32 {
|
|
find_part_numbers(input)
|
|
.filter_map(|(num, mut parts)| parts.next().map(|_| num))
|
|
.sum()
|
|
}
|
|
|
|
fn part2(input: &Input) -> u32 {
|
|
let mut gears = FxHashMap::<(usize, usize), _>::default();
|
|
for (num, gs) in find_part_numbers(input) {
|
|
for g in gs {
|
|
gears.entry(g).or_insert_with(Vec::new).push(num);
|
|
}
|
|
}
|
|
gears
|
|
.into_iter()
|
|
.filter(|&((i, j), ref nums)| {
|
|
matches!(input[i][j], Cell::Symbol { gear: true }) && nums.len() == 2
|
|
})
|
|
.map(|(_, nums)| nums.into_iter().product::<u32>())
|
|
.sum()
|
|
}
|
|
|
|
aoc::main!(2023, 3, ex: 1);
|