adventofcode2022/day12/common.rs

// Copyright 2022 Christian Ulrich
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use std::collections::HashSet;
use std::io;

pub struct Map {
    squares: Vec<Vec<u8>>,
    start_pos: (usize, usize),
    end_pos: (usize, usize),
}

#[derive(Debug, Clone)]
struct DijkstraInfo {
    pos: (usize, usize),
    shortest_distance: usize,
    previous: Option<(usize, usize)>,
}

#[derive(Debug)]
pub struct ShortestPath {
    dijkstra_table: Vec<DijkstraInfo>,
    current: (usize, usize),
    start_pos: (usize, usize),
}

impl Map {
    pub fn from_stdin() -> Result<Self, &'static str> {
        const ERROR: &str = "cannot parse map";
        let squares: Result<Vec<Vec<u8>>, _> = io::stdin()
            .lines()
            .map(|line| Ok(line.map_err(|_| ERROR)?.into_bytes()))
            .collect();
        let mut squares = squares?;
        let mut start = None;
        let mut end = None;
        for y in 0..squares.len() {
            if squares[y].len() != squares[0].len() {
                return Err(ERROR);
            }
            for x in 0..squares[y].len() {
                match squares[y][x] as char {
                    'S' if { start.is_none() } => {
                        start = Some((x, y));
                        squares[y][x] = 'a' as u8;
                    }
                    'S' if { start.is_some() } => return Err(ERROR),
                    _ => (),
                }
                match squares[y][x] as char {
                    'E' if { end.is_none() } => {
                        end = Some((x, y));
                        squares[y][x] = 'z' as u8;
                    }
                    'E' if { end.is_some() } => return Err(ERROR),
                    _ => (),
                }
            }
        }
        Ok(Self {
            squares: squares,
            start_pos: start.ok_or(ERROR)?,
            end_pos: end.ok_or(ERROR)?,
        })
    }

    pub fn square_height(&self, pos: (usize, usize)) -> Option<u8> {
        let (x, y) = pos;
        self.squares.get(y)?.get(x).cloned()
    }

    pub fn start_pos(&self) -> (usize, usize) {
        self.start_pos
    }

    pub fn end_pos(&self) -> (usize, usize) {
        self.end_pos
    }

    pub fn neighbors(&self, pos: (usize, usize)) -> Option<Vec<(usize, usize)>> {
        let mut result = vec![];
        let height = self.square_height(pos)?;
        let (x, y) = pos;
        if x > 0 {
            result.push((x - 1, y));
        }
        if x < self.squares[0].len() - 1 {
            result.push((x + 1, y));
        }
        if y > 0 {
            result.push((x, y - 1));
        }
        if y < self.squares.len() - 1 {
            result.push((x, y + 1));
        }
        let is_accessible = |pos: (usize, usize)| {
            let neighbor_height = self.square_height(pos).unwrap();
            u8::abs_diff(neighbor_height, height) <= 1 || neighbor_height > height
        };
        Some(
            result
                .iter()
                .filter(|&pos| is_accessible(*pos))
                .cloned()
                .collect(),
        )
    }

    fn dijkstra_table(&self, start_pos: (usize, usize)) -> Vec<DijkstraInfo> {
        let mut result = Vec::with_capacity(self.squares.len() * self.squares[0].len());
        for y in 0..self.squares.len() {
            for x in 0..self.squares[0].len() {
                let info = DijkstraInfo {
                    pos: (x, y),
                    shortest_distance: if (x, y) == start_pos { 0 } else { usize::MAX },
                    previous: None,
                };
                result.push(info);
            }
        }
        result
    }

    pub fn shortest_path(&self, start_pos: Option<(usize, usize)>) -> ShortestPath {
        // we reverse start_pos and end_pos so the ShortestPath will have the right order
        let end_pos = match start_pos {
            Some(pos) => pos,
            None => self.start_pos,
        };
        let start_pos = self.end_pos;

        let mut table = self.dijkstra_table(start_pos);
        let mut unvisited: HashSet<(usize, usize)> = table.iter().map(|i| i.pos).collect();
        while let Some(info) = table
            .iter()
            .filter(|i| i.shortest_distance < usize::MAX && unvisited.contains(&i.pos))
            .min_by_key(|i| i.shortest_distance)
            .cloned()
        {
            for neighbor_pos in self.neighbors(info.pos).unwrap() {
                let neighbor_info = table.iter_mut().find(|n| n.pos == neighbor_pos).unwrap();
                let distance = info.shortest_distance + 1;
                if distance < neighbor_info.shortest_distance {
                    neighbor_info.shortest_distance = distance;
                    neighbor_info.previous = Some(info.pos);
                }
            }
            unvisited.remove(&info.pos);
        }

        ShortestPath {
            dijkstra_table: table,
            current: end_pos,
            start_pos: start_pos,
        }
    }

    pub fn lowest_positions(&self) -> Vec<(usize, usize)> {
        let mut result = vec![];
        for y in 0..self.squares.len() {
            for x in 0..self.squares[0].len() {
                if self.squares[y][x] as char == 'a' {
                    result.push((x, y));
                }
            }
        }
        result
    }
}

impl Iterator for ShortestPath {
    type Item = (usize, usize);

    fn next(&mut self) -> Option<Self::Item> {
        if self.current == self.start_pos {
            return None;
        }
        let result = self.current;
        let info = self
            .dijkstra_table
            .iter()
            .find(|i| i.pos == self.current)
            .unwrap();
        self.current = info.previous?;
        Some(result)
    }
}