AoC2021/src/day15.rs

use std::cmp::{max, Reverse};
use std::collections::{BinaryHeap, BTreeMap, HashMap};
use crate::Day;

type Coordinate = (usize, usize);

type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;

pub struct Day15(HashMap<Coordinate, u32>, usize, usize);

impl Day for Day15 {
    fn init(content: String) -> anyhow::Result<Self> {
        let mut m = HashMap::new();
        let mut y_max = 0;
        let mut x_max = 0;
        for (y, line) in content.lines().enumerate() {
            for (x, ch) in line.chars().enumerate() {
                m.insert((x, y), ch.to_digit(10).ok_or_else(|| anyhow::Error::msg("Could not parse digit as number"))?);
                x_max = max(x, x_max);
            }
            y_max = max(y, y_max);
        }

        Ok(Self(m, x_max, y_max))
    }

    fn part1(&self) -> anyhow::Result<String> {
        let g = map_to_graph(&self.0, (self.1, self.1));
        let end = &(self.1, self.2);
        let r = dijkstra(&g, &(0, 0), end).unwrap();
        Ok(format!("{}", r.1))
    }

    fn part2(&self) -> anyhow::Result<String> {
        let mut m = HashMap::new();
        let len = self.1 + 1;
        for y in 0..len {
            for x in 0..len {
                let og = (x, y);
                let orig_value = *self.0.get(&og).unwrap();
                for ry in 0..5 {
                    for rx in 0..5 {
                        let coord = (rx * len + x, ry * len + y);
                        let e = m.entry(coord).or_insert(0);
                        *e = orig_value + rx as u32 + ry as u32;
                        if *e > 9 {
                            *e -= 9;
                        }
                    }
                }
            }
        }

        let target = (len * 5 - 1, len * 5 - 1);
        let g = map_to_graph(&m, target);
        let r = dijkstra(&g, &(0, 0), &target).unwrap();
        Ok(format!("{}", r.1))
    }
}

fn neighbours((x, y): Coordinate, (tx, ty): Coordinate) -> impl Iterator<Item=Coordinate> {
    [
        x.checked_sub(1).map(|r| (r, y)),
        y.checked_sub(1).map(|c| (x, c)),
        (x < tx).then(|| (x + 1, y)),
        (y < ty).then(|| (x, y + 1)),
    ].into_iter().flatten()
}

pub fn dijkstra(
    graph: &Graph<Coordinate, u32>,
    start: &Coordinate,
    end: &Coordinate,
) -> Option<(Coordinate, u32)> {
    let mut ans = BTreeMap::new();
    let mut prio = BinaryHeap::new();

    // start is the special case that doesn't have a predecessor
    ans.insert(*start, None);

    for (new, weight) in &graph[start] {
        ans.insert(*new, Some((*start, *weight)));
        prio.push(Reverse((*weight, new, start)));
    }

    while let Some(Reverse((dist_new, new, prev))) = prio.pop() {
        match ans[new] {
            // what we popped is what is in ans, we'll compute it
            Some((p, d)) if p == *prev && d == dist_new => {}
            // otherwise it's not interesting
            _ => continue,
        }

        for (next, weight) in &graph[new] {
            match ans.get(next) {
                // if ans[next] is a lower dist than the alternative one, we do nothing
                Some(Some((_, dist_next))) if dist_new + *weight >= *dist_next => {}
                // if ans[next] is None then next is start and so the distance won't be changed, it won't be added again in prio
                Some(None) => {}
                // the new path is shorter, either new was not in ans or it was farther
                _ => {
                    ans.insert(*next, Some((*new, *weight + dist_new)));
                    prio.push(Reverse((*weight + dist_new, next, new)));
                }
            }
        }
    }

    *ans.get(end).unwrap()
}

fn add_edge(graph: &mut Graph<Coordinate, u32>, source: Coordinate, target: Coordinate, value: u32) {
    graph.entry(source).or_insert_with(BTreeMap::new).insert(target, value);
    graph.entry(target).or_insert_with(BTreeMap::new);
}

fn map_to_graph(map: &HashMap<Coordinate, u32>, end: Coordinate) -> Graph<Coordinate, u32> {
    let mut g = BTreeMap::new();
    for y in 0..=end.1 {
        for x in 0..=end.0 {
            for neighbour in neighbours((x, y), end) {
                add_edge(&mut g, (x, y), neighbour, *map.get(&neighbour).unwrap())
            }
        }
    }
    g
}

#[cfg(test)]
mod tests {
    use crate::Day;
    use crate::day15::Day15;

    const INPUT: &str = r"1163751742
1381373672
2136511328
3694931569
7463417111
1319128137
1359912421
3125421639
1293138521
2311944581";

    #[test]
    fn part1_test() -> anyhow::Result<()> {
        let d = Day15::init(INPUT.to_string())?;
        assert_eq!("40", d.part1()?);
        Ok(())
    }

    #[test]
    fn part2_test() -> anyhow::Result<()> {
        let d = Day15::init(INPUT.to_string())?;
        assert_eq!("315", d.part2()?);
        Ok(())
    }

    // Woo slow and i again don't feel like fixing it :p
    // #[test]
    // fn part1_real() -> anyhow::Result<()> {
    //     let d = Day15::init(crate::load_input("15")?)?;
    //     assert_eq!("388", d.part1()?);
    //     Ok(())
    // }
    //
    // #[test]
    // fn part2_real() -> anyhow::Result<()> {
    //     let d = Day15::init(crate::load_input("15")?)?;
    //     assert_eq!("2819", d.part2()?);
    //     Ok(())
    // }
}