/ Rust / 2023 / 08.rs
08.rs
  1  #![feature(test)]
  2  
  3  use itertools::Itertools;
  4  use num::Integer;
  5  use rayon::prelude::*;
  6  use regex::Regex;
  7  use rustc_hash::FxHashMap;
  8  
  9  #[derive(Debug)]
 10  struct Input {
 11      instructions: Vec<Instruction>,
 12      nodes: Vec<Node>,
 13      start: usize,
 14      goal: usize,
 15  }
 16  
 17  #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 18  enum Instruction {
 19      Left,
 20      Right,
 21  }
 22  
 23  #[derive(Debug, Clone, Copy)]
 24  struct Node {
 25      left: usize,
 26      right: usize,
 27      start: bool,
 28      goal: bool,
 29  }
 30  
 31  fn setup(input: &str) -> Input {
 32      let instructions = input
 33          .lines()
 34          .next()
 35          .unwrap()
 36          .bytes()
 37          .map(|b| match b {
 38              b'L' => Instruction::Left,
 39              b'R' => Instruction::Right,
 40              _ => panic!(),
 41          })
 42          .collect();
 43  
 44      let names = input
 45          .lines()
 46          .skip(2)
 47          .enumerate()
 48          .map(|(i, line)| (line.split_whitespace().next().unwrap(), i))
 49          .collect::<FxHashMap<_, _>>();
 50  
 51      let regex = Regex::new(r"([^, ]+) = \(([^, ]+), ([^, ]+)\)").unwrap();
 52      let nodes = input
 53          .lines()
 54          .skip(2)
 55          .map(|line| {
 56              let captures = regex.captures(line).unwrap();
 57              let left = names[&captures[2]];
 58              let right = names[&captures[3]];
 59              Node {
 60                  left,
 61                  right,
 62                  start: captures[1].ends_with('A'),
 63                  goal: captures[1].ends_with('Z'),
 64              }
 65          })
 66          .collect();
 67  
 68      Input {
 69          instructions,
 70          nodes,
 71          start: names.get("AAA").copied().unwrap_or_default(),
 72          goal: names.get("ZZZ").copied().unwrap_or_default(),
 73      }
 74  }
 75  
 76  fn solve(input: &Input, start: usize, is_goal: impl Fn(usize) -> bool) -> usize {
 77      Itertools::take_while_inclusive(
 78          input
 79              .instructions
 80              .iter()
 81              .cycle()
 82              .scan(start, |p, instruction| {
 83                  let node = input.nodes[*p];
 84                  *p = match instruction {
 85                      Instruction::Left => node.left,
 86                      Instruction::Right => node.right,
 87                  };
 88                  Some(*p)
 89              }),
 90          |&p| !is_goal(p),
 91      )
 92      .count()
 93  }
 94  
 95  fn part1(input: &Input) -> usize {
 96      solve(input, input.start, |p| p == input.goal)
 97  }
 98  
 99  fn part2(input: &Input) -> usize {
100      input
101          .nodes
102          .par_iter()
103          .enumerate()
104          .filter(|(_, node)| node.start)
105          .map(|(i, _)| solve(input, i, |p| input.nodes[p].goal))
106          .reduce(|| 1, |a, b| a.lcm(&b))
107  }
108  
109  aoc::main!(2023, 8, ex: 1[a], 2[a], 3[b]);