path: root/src/day_17.rs

use std::{
    cmp::Reverse,
    collections::{hash_map::Entry, BinaryHeap, HashMap},
};

use anyhow::Context;

use crate::{Problem, Solution};

pub struct Day17;

impl Problem for Day17 {
    const DAY: u8 = 17;

    const INPUT: &'static str = include_str!("../input/day_17.txt");
}

impl Solution for Day17 {
    type Answer1 = usize;

    type Answer2 = usize;

    fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
        let grid = parse_grid(input).context("Failed to parse grid")?;

        astar::<1, 3>(&grid, (0, 0), (grid.len() - 1, grid[0].len() - 1))
            .context("Failed to find path")

        // 755 ==
    }

    fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
        let grid = parse_grid(input).context("Failed to parse grid")?;

        astar::<4, 10>(&grid, (0, 0), (grid.len() - 1, grid[0].len() - 1))
            .context("Failed to find path")
        // 879 <
        // 881 ==
    }
}

fn parse_grid(input: &str) -> Option<Grid> {
    input
        .lines()
        .map(|v| {
            v.chars()
                .map(|c| c.to_digit(10).map(|n| n as usize))
                .try_collect::<Vec<_>>()
        })
        .try_collect::<Grid>()
}

#[cfg(test)]
fn rebuild_path(edges: &[Edge], start: Position, end: Position) -> Vec<Position> {
    let mut path = vec![];
    let mut position = start;
    for edge in edges {
        path.push(position);
        position = edge.traverse(position).unwrap();
    }
    path.push(end);
    path
}

fn astar<const S: usize, const M: usize>(
    grid: &Grid,
    start: Position,
    goal: Position,
) -> Option<usize> {
    let mut queue = BinaryHeap::from([Reverse(Node::<S, M>::start(start))]);
    let mut cache = HashMap::from([((start, None), usize::MAX)]);

    #[cfg(test)]
    let printer = &mut crate::Printer::from(std::io::stdout().lock());

    while let Some(Reverse(node)) = queue.pop() {
        for (edges, position, cost) in node.successors(grid) {
            let mut next = node.clone();
            next.position = position;
            next.path.extend(edges);
            next.cost += cost;
            next.estimated = next.cost + next.dist(goal);

            #[cfg(test)]
            {
                let node = &next;
                let path = rebuild_path(&node.path, start, node.position);
                printer
                    .with(format_args!("pos: {:?}", node.position))
                    .with(format_args!("score: {}", node.cost))
                    .with(format_args!("trail: {:?}", node.trail()))
                    .with_grid(grid, &path)
                    .clear()
                    .write()
                    .unwrap();
            }

            // pause();
            if next.position == goal {
                #[cfg(test)]
                printer
                    .with(format_args!("score: {}", node.cost))
                    .with_grid(grid, &rebuild_path(&next.path, start, next.position))
                    .clear()
                    .write()
                    .unwrap();
                return Some(next.cost);
            }

            match cache.entry((next.position, next.trail())) {
                Entry::Vacant(e) => {
                    e.insert(next.cost);
                }
                Entry::Occupied(mut e) => {
                    if next.cost < *e.get() {
                        e.insert(next.cost);
                    } else {
                        continue;
                    }
                }
            }

            queue.push(Reverse(next))
        }
    }

    None
}

type Grid = Vec<Vec<usize>>;
type Position = (usize, usize);

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct Node<const STEP: usize, const MAX: usize> {
    estimated: usize,
    cost: usize,
    position: Position,
    path: Vec<Edge>,
}

impl<const S: usize, const M: usize> Node<S, M> {
    fn start(position: Position) -> Self {
        Node {
            position,
            estimated: usize::MAX,
            ..Default::default()
        }
    }

    fn trail(&self) -> Option<[Edge; M]> {
        self.path.last_chunk().copied()
    }

    fn successors(&self, grid: &Grid) -> Vec<(Vec<Edge>, Position, usize)> {
        let (iter, last_seq) = match self.path.last() {
            Some(last) => (
                last.edges(),
                Some((
                    last,
                    self.path.iter().rev().take_while(|e| *e == last).count(),
                )),
            ),
            None => (Edge::ALL.as_slice(), None),
        };

        iter.iter()
            .flat_map(|edge| {
                match last_seq {
                    Some((last, seq)) if last == edge => 1..=(M - seq),
                    _ => S..=M,
                }
                .map(move |step| (edge, step))
            })
            .filter_map(|(edge, step)| self.traverse(*edge, step, grid))
            .collect()
    }

    fn traverse(
        &self,
        edge: Edge,
        step: usize,
        grid: &Grid,
    ) -> Option<(Vec<Edge>, Position, usize)> {
        let mut cost = 0;
        let mut position = self.position;

        for _ in 0..step {
            position = edge.traverse(position)?;
            cost += grid.get(position.0)?.get(position.1)?;
        }

        Some((vec![edge; step], position, cost))
    }

    fn dist(&self, p2: Position) -> usize {
        (self.position.0.abs_diff(p2.0) + self.position.1.abs_diff(p2.1)) * S
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum Edge {
    North,
    East,
    South,
    West,
}

impl Edge {
    const ALL: [Edge; 4] = [Edge::North, Edge::East, Edge::South, Edge::West];

    fn edges(&self) -> &[Edge] {
        match self {
            Edge::North => &[Edge::North, Edge::East, Edge::West],
            Edge::East => &[Edge::North, Edge::East, Edge::South],
            Edge::South => &[Edge::East, Self::South, Edge::West],
            Edge::West => &[Edge::North, Edge::South, Edge::West],
        }
    }

    fn traverse(&self, (row, col): Position) -> Option<Position> {
        Some(match self {
            Edge::North => (row.checked_sub(1)?, col),
            Edge::East => (row, col.checked_add(1)?),
            Edge::South => (row.checked_add(1)?, col),
            Edge::West => (row, col.checked_sub(1)?),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const INPUT: &str = indoc::indoc! {r#"
        2413432311323
        3215453535623
        3255245654254
        3446585845452
        4546657867536
        1438598798454
        4457876987766
        3637877979653
        4654967986887
        4564679986453
        1224686865563
        2546548887735
        4322674655533
    "#};

    const INPUT2: &str = indoc::indoc! {r#"
        111111111111
        999999999991
        999999999991
        999999999991
        999999999991
    "#};

    #[test]
    fn test_part_1() -> anyhow::Result<()> {
        Ok(assert_eq!(102, Day17::part_1(INPUT)?))
    }

    #[test]
    fn test_part_2_1() -> anyhow::Result<()> {
        Ok(assert_eq!(94, Day17::part_2(INPUT)?))
    }

    #[test]
    fn test_part_2_2() -> anyhow::Result<()> {
        Ok(assert_eq!(71, Day17::part_2(INPUT2)?))
    }
}