advent-of-code/2023/day8/rust/src/main.rs

#![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<Step>,
}

impl<'a> StepsWalker<'a> {
    pub fn new(start: &'a str, steps: Vec<Step>) -> 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<G>(&mut self, context: G) -> Option<&'a str>
    where
        G: IntoEdgesDirected + GraphBase<NodeId = &'a str> + Data<EdgeWeight = Edge>,
    {
        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<G>(mut self, context: G) -> impl Iterator<Item = usize>
    where
        G: IntoEdgesDirected + GraphBase<NodeId = &'a str> + Data<EdgeWeight = Edge>,
    {
        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::<Vec<_>, _>(|&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<G> for StepsWalker<'a>
where
    G: IntoEdgesDirected + GraphBase<NodeId = &'a str> + Data<EdgeWeight = Edge>,
{
    type Item = &'a str;

    fn walk_next(&mut self, context: G) -> Option<Self::Item> {
        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}");
}