2022/src/bin/day09.rs

use std::{collections::BTreeSet, str::FromStr};

use aoc_2022::*;

use color_eyre::{eyre::eyre, Report};
use itertools::Itertools;

type Input = Vec<Instruction>;
type Output = usize;

enum Direction {
    Left,
    Right,
    Down,
    Up,
}

impl Direction {
    fn apply(&self, pos: Position) -> Position {
        pos + match self {
            Direction::Left => Position::new(-1, 0),
            Direction::Right => Position::new(1, 0),
            Direction::Down => Position::new(0, -1),
            Direction::Up => Position::new(0, 1),
        }
    }
}

impl FromStr for Direction {
    type Err = Report;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "L" => Ok(Direction::Left),
            "R" => Ok(Direction::Right),
            "D" => Ok(Direction::Down),
            "U" => Ok(Direction::Up),
            _ => Err(eyre!("Invalid direction {}", s)),
        }
    }
}

struct Instruction {
    direction: Direction,
    steps: i32,
}

impl Instruction {
    #[cfg(test)]
    fn new(direction: Direction, steps: i32) -> Self {
        Self { direction, steps }
    }
}

impl FromStr for Instruction {
    type Err = Report;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let parts = s.split(' ').collect_vec();
        let (direction, steps) = (parts[0].parse::<Direction>()?, parts[1].parse::<i32>()?);

        Ok(Self { direction, steps })
    }
}

type Position = Vector2D<i32>;

#[derive(Debug, Clone)]
struct State {
    knots: Vec<Position>,
}

impl State {
    fn new(knots: i32) -> State {
        Self {
            knots: (0..knots).map(|_| Position::new(0, 0)).collect_vec(),
        }
    }

    fn execute_step(&self, visited: &mut BTreeSet<Position>, direction: &Direction) -> State {
        let mut knots = self.knots.clone();

        // move head
        knots[0] = direction.apply(knots[0]);

        // index of tail
        let tail = knots.len() - 1;

        // add old position of the tail to the set
        visited.insert(knots[tail]);

        // move parts towards head
        for i in 1..knots.len() {
            let v = knots[i - 1] - knots[i];
            if v.x().abs() > 1 || v.y().abs() > 1 {
                let d = Position::new(v.x().signum(), v.y().signum());
                knots[i] = knots[i] + d;
            }
        }
        // debug!("\n{}", State { head, tail }.show(6, 5));

        // add new position of the tail to the set
        visited.insert(knots[tail]);

        State { knots }
    }

    fn execute(&self, visited: &mut BTreeSet<Position>, i: &Instruction) -> State {
        // debug!("\n{}", self.show(6, 5));
        let mut state: State = self.clone();
        for _ in 0..i.steps {
            state = state.execute_step(visited, &i.direction);
        }

        state
    }

    // fn show(&self, width: i32, height: i32) -> String {
    //     (0..height)
    //         .rev()
    //         .map(|y| {
    //             (0..width)
    //                 .map(|x| {
    //                     let pos = Position::new(x, y);

    //                     if pos == self.head {
    //                         'H'
    //                     } else if pos == self.tail {
    //                         'T'
    //                     } else {
    //                         '.'
    //                     }
    //                 })
    //                 .join("")
    //         })
    //         .join("\n")
    // }
}

fn execute(knots: i32, input: &Input) -> Output {
    let mut visited = BTreeSet::<Position>::new();

    input.iter().fold(State::new(knots), |state, instruction| {
        state.execute(&mut visited, instruction)
    });

    visited.len()
}

struct Day09;
impl Solution<Input, Output> for Day09 {
    fn parse_input<P: AsRef<Path>>(pathname: P) -> Input {
        file_to_structs(pathname)
    }

    fn part_1(input: &Input) -> Output {
        execute(2, input)
    }

    fn part_2(input: &Input) -> Output {
        execute(10, input)
    }
}

fn main() -> Result<()> {
    Day09::main()
}

test_sample!(day_09, Day09, 13, 1);

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

    #[test]
    fn test_part_2_bigger_sample() {
        let input = vec![
            Instruction::new(Direction::Right, 5),
            Instruction::new(Direction::Up, 8),
            Instruction::new(Direction::Left, 8),
            Instruction::new(Direction::Down, 3),
            Instruction::new(Direction::Right, 17),
            Instruction::new(Direction::Down, 10),
            Instruction::new(Direction::Left, 25),
            Instruction::new(Direction::Up, 20),
        ];

        assert_eq!(Day09::part_2(&input), 36);
    }
}