AoC/src/day6.rs

type Input = ((usize, usize), Vec<Vec<u8>>);

use aoc_runner_derive::{aoc, aoc_generator};
#[aoc_generator(day6)]
fn parse(input: &str) -> Input {
    let lines = input.split("\n");
    let mut pos = (0, 0);
    let map = lines
        .enumerate()
        .map(|(y, line)| {
            line.chars()
                .enumerate()
                .map(|(x, c)| match c {
                    '.' => 0,
                    '#' => 1,
                    '^' => {
                        pos = (x, y);
                        0
                    }
                    _ => todo!(),
                })
                .collect()
        })
        .collect();
    (pos, map)
}

#[aoc(day6, part1)]
fn part1(input: &Input) -> u64 {
    let h = input.1.len();
    let w = input.1[0].len();
    let mut map = input.1.clone();
    let mut pos = input.0;
    let mut dir = (0, -1);
    let mut count = 0;
    loop {
        if map[pos.1][pos.0] == 0 {
            map[pos.1][pos.0] = 2;
            count += 1;
        };
        if pos.0 as isize + dir.0 < 0
            || pos.0 as isize + dir.0 >= w as isize
            || pos.1 as isize + dir.1 < 0
            || pos.1 as isize + dir.1 >= h as isize
        {
            break;
        }
        while map[(pos.1 as isize + dir.1) as usize][(pos.0 as isize + dir.0) as usize] == 1 {
            if dir.0 == 1 {
                dir = (0, 1);
            } else if dir.0 == -1 {
                dir = (0, -1);
            } else if dir.1 == 1 {
                dir = (-1, 0);
            } else if dir.1 == -1 {
                dir = (1, 0);
            }
        }
        pos = (
            (pos.0 as isize + dir.0) as usize,
            (pos.1 as isize + dir.1) as usize,
        )
    }
    count
}

#[aoc(day6, part2)]
fn part2(input: &Input) -> u64 {
    let initial = input.0;
    let h = input.1.len();
    let w = input.1[0].len();
    let mut map = input.1.clone();
    let mut pos = input.0;
    let mut dir = (0, -1);
    let mut loops = vec![];
    loop {
        if pos.0 as isize + dir.0 < 0
            || pos.0 as isize + dir.0 >= w as isize
            || pos.1 as isize + dir.1 < 0
            || pos.1 as isize + dir.1 >= h as isize
        {
            break;
        }
        while map[(pos.1 as isize + dir.1) as usize][(pos.0 as isize + dir.0) as usize] == 1 {
            if dir.0 == 1 {
                dir = (0, 1);
            } else if dir.0 == -1 {
                dir = (0, -1);
            } else if dir.1 == 1 {
                dir = (-1, 0);
            } else if dir.1 == -1 {
                dir = (1, 0);
            }
        }
        let next_pos = (
            (pos.0 as isize + dir.0) as usize,
            (pos.1 as isize + dir.1) as usize,
        );
        if map[next_pos.1][next_pos.0] != 1 {
            map[next_pos.1][next_pos.0] = 1;
            if check_loop(initial, (0, -1), &map) {
                // count += 1;
                if !(loops.contains(&next_pos)) && next_pos != initial {
                    loops.push(next_pos);
                }
                // println!("{} {}", nexts_pos.0 + 1, next_pos.1 + 1);
            }
            map[next_pos.1][next_pos.0] = 0;
        }
        pos = next_pos;
    }
    loops.len() as u64
}

fn check_loop(pos: (usize, usize), dir: (isize, isize), map: &[Vec<u8>]) -> bool {
    let mut visited = vec![];

    let h = map.len();
    let w = map[0].len();

    let mut dir = dir;
    let mut pos = pos;

    loop {
        if pos.0 as isize + dir.0 < 0
            || pos.0 as isize + dir.0 >= w as isize
            || pos.1 as isize + dir.1 < 0
            || pos.1 as isize + dir.1 >= h as isize
        {
            break;
        }
        let mut turns = 0;
        while map[(pos.1 as isize + dir.1) as usize][(pos.0 as isize + dir.0) as usize] == 1 {
            if dir.0 == 1 {
                dir = (0, 1);
            } else if dir.0 == -1 {
                dir = (0, -1);
            } else if dir.1 == 1 {
                dir = (-1, 0);
            } else if dir.1 == -1 {
                dir = (1, 0);
            }

            turns += 1;
        }
        pos = (
            (pos.0 as isize + dir.0) as usize,
            (pos.1 as isize + dir.1) as usize,
        );
        if visited.contains(&(pos.0, pos.1, dir.0, dir.1)) {
            return true;
        } else if turns > 0 {
            visited.push((pos.0, pos.1, dir.0, dir.1));
        }
    }

    false
}

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

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse("....#.....\n.........#\n..........\n..#.......\n.......#..\n..........\n.#..^.....\n........#.\n#.........\n......#...")), 41);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse("....#.....\n.........#\n..........\n..#.......\n.......#..\n..........\n.#..^.....\n........#.\n#.........\n......#...")), 6);
    }
}