path: root/2023/rs/code/three/src/main.rs

#![feature(iter_map_windows)]

use std::collections::BTreeSet;

static INPUT: &[u8] = include_bytes!("input.txt");

#[derive(Copy, Clone, Debug)]
enum LineMaskEntry {
  Empty,
  Gear,
  Number{tag: u32, num: u32}
}

fn main() {
  let valueified = INPUT.split(|b| *b == b'\n').filter(|s| !s.is_empty());

  let mut num = 0;

  let (parts_number_sum, gears_mask) = [valueified.clone().next().unwrap()].into_iter().chain(valueified.clone()).chain([valueified.clone().last().unwrap()])
    .map_windows::<_, _, 3>(|[prev, curr, next]| {
      let mut numlen = 0;
      let mut numval = 0;

      let mut acc = 0;
      let mut linemask = vec![LineMaskEntry::Empty; curr.len()];

      for i in 0..curr.len() {
        if curr[i] == b'*' { linemask[i] = LineMaskEntry::Gear; }

        if !curr[i].is_ascii_digit() && numlen != 0 {
          let start = if i - numlen == 0 { i - numlen } else { i - numlen - 1 };
          let end = i;
          let window = || start..=end;
          if prev[window()].into_iter().chain(&curr[window()]).chain(&next[window()]).any(|b| !b.is_ascii_digit() && *b != b'.') {
            acc += numval;

            linemask[i - numlen .. i].fill(LineMaskEntry::Number{ tag: num, num: numval });
            num += 1;
          }

          numlen = 0;
          numval = 0;
        }
        else if curr[i].is_ascii_digit() {
          numlen += 1;
          numval *= 10;
          numval += (curr[i] - b'0') as u32;
        }
      }

      if numlen != 0 {
        let window = || curr.len() - numlen - 1..curr.len();
        if prev[window()].into_iter().chain(&curr[window()]).chain(&next[window()]).any(|b| !b.is_ascii_digit() && *b != b'.') {
          acc += numval;

          linemask[curr.len() - numlen .. curr.len()].fill(LineMaskEntry::Number{ tag: num, num: numval });
          num += 1;
        }
      }
      (acc, linemask)
    }).fold((0, Vec::new()), |(n, mut v), (a, l)| (n + a, {v.push(l); v}));
  println!("Sum of valid part numbers: {parts_number_sum}");


  let first = vec![LineMaskEntry::Empty; gears_mask[0].len()];
  let last = first.clone();
  let gear_ratio_sum = [first].into_iter().chain(gears_mask).chain([last])
    .map_windows::<_, _, 3>(|[prev, curr, next]| {
      let mut acc = 0;

      for i in 0..curr.len() {
        match curr[i] {
          LineMaskEntry::Gear => {
            let mut set = BTreeSet::new();
            for row in [prev, curr, next] {
              for n in [i - 1, i, i + 1] {
                if n == 0 || n == row.len() { continue; }
                if let LineMaskEntry::Number{tag, num} = row[n] {
                  set.insert((tag, num));
                }
              }
            }
            if set.len() == 2 {
              let mut iter = set.into_iter();
              let x = iter.next().unwrap().1;
              let y = iter.next().unwrap().1;
              acc += x * y
            }
          },
          _ => ()
        }
      }
      acc
    }).sum::<u32>();
  println!("Gear ratio sum: {gear_ratio_sum}");
}