#![warn(clippy::pedantic)] use std::{cmp::PartialEq, collections::HashMap, fmt::Debug, io::stdin, num::NonZeroUsize}; use either::Either; use itertools::Itertools; use petgraph::{ prelude::DiGraphMap, visit::{Data, EdgeRef, GraphBase, IntoEdgesDirected, Walker}, Direction::Outgoing, }; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum Step { Left, Right, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum Edge { Left, Right, Both, } impl Edge { pub fn matches(self, s: Step) -> bool { match (self, s) { (Edge::Left, Step::Left) | (Edge::Right, Step::Right) | (Edge::Both, _) => true, (Edge::Right, Step::Left) | (Edge::Left, Step::Right) => false, } } } struct StepsWalker<'a> { current_node: &'a str, step_idx: usize, steps: Vec, } impl<'a> StepsWalker<'a> { pub fn new(start: &'a str, steps: Vec) -> Self { Self { current_node: start, steps, step_idx: 0, } } /// Resets the step sequence to the beginning. pub fn reset_steps(&mut self) { self.step_idx = 0; } /// Steps once. Returns `None` if all steps have been performed. pub fn step(&mut self, context: G) -> Option<&'a str> where G: IntoEdgesDirected + GraphBase + Data, { let step = self.steps.get(self.step_idx)?; let mut targets = context .edges_directed(self.current_node, Outgoing) .filter_map(|e| e.weight().matches(*step).then_some(e.target())); let next_node = targets.next()?; assert_eq!( targets.next(), None, "More than one {:?} edge from {:?}", step, self.current_node ); self.current_node = next_node; self.step_idx += 1; Some(self.current_node) } pub fn find_end_positions(mut self, context: G) -> impl Iterator where G: IntoEdgesDirected + GraphBase + Data, { let mut seen = HashMap::new(); let mut pos = 1_usize; loop { while let Some(node) = self.step(context) { let key = (node, self.step_idx - 1); if let Some(&loop_start) = seen.get(&key) { let period = NonZeroUsize::new(pos - loop_start).unwrap(); let mut ends: Vec<_> = seen .into_iter() .filter_map(|((node, _step_idx), pos)| node.ends_with('Z').then_some(pos)) .collect(); ends.sort_unstable(); let (initial, repeating) = ends .into_iter() .partition::, _>(|&pos| pos < loop_start); return initial.into_iter().chain((0_usize..).flat_map(move |n| { repeating .clone() .into_iter() .map(move |pos| pos + n * period.get()) })); } seen.insert(key, pos); pos += 1; } self.reset_steps(); } } } impl<'a, G> Walker for StepsWalker<'a> where G: IntoEdgesDirected + GraphBase + Data, { type Item = &'a str; fn walk_next(&mut self, context: G) -> Option { Some(if let Some(node) = self.step(context) { node } else { self.reset_steps(); self.step(context).unwrap() }) } } fn main() { let mut lines = stdin().lines().map(Result::unwrap); let steps: Vec<_> = lines .next() .unwrap() .chars() .map(|c| match c { 'L' => Step::Left, 'R' => Step::Right, _ => unreachable!(), }) .collect(); let lines: Vec<_> = lines.skip(1).collect(); let graph: DiGraphMap<_, _> = lines .iter() .flat_map(|l| { let (from, to) = l.split_once(" = ").unwrap(); let (left, right) = to .strip_prefix('(') .unwrap() .strip_suffix(')') .unwrap() .split_once(", ") .unwrap(); if left == right { Either::Left([(from, left, Edge::Both)]) } else { Either::Right([(from, left, Edge::Left), (from, right, Edge::Right)]) } .into_iter() }) .collect(); let walker = StepsWalker::new("AAA", steps.clone()); println!( "{}", walker.iter(&graph).position(|n| n == "ZZZ").unwrap() + 1 ); let start_nodes = graph.nodes().filter(|n| n.ends_with('A')); let mut end_positions: Vec<_> = start_nodes .map(|start| { StepsWalker::new(start, steps.clone()) .find_end_positions(&graph) .peekable() }) .collect(); let num_steps = loop { let target = *end_positions[0].peek().unwrap(); for positions in &mut end_positions { positions .peeking_take_while(|&pos| pos < target) .for_each(|_| {}); } if end_positions .iter_mut() .map(|positions| positions.peek()) .all_equal() { break target; } end_positions[0].next().unwrap(); }; println!("{num_steps}"); }