AoC2021/src/day12.rs

use std::collections::HashMap;
use crate::Day;

pub struct Day12(HashMap<String, Cave>);

impl Day for Day12 {
    fn init(content: String) -> anyhow::Result<Self> {
        let mut caves = HashMap::new();

        for line in content.lines() {
            let z = line.split('-').collect::<Vec<_>>();
            let first = z[0];
            let second = z[1];

            {
                let first_cave = caves.entry(first.to_string()).or_insert_with(|| Cave::new(first.into()));
                first_cave.add_connection(second);
            }
            {
                let second_cave = caves.entry(second.to_string()).or_insert_with(|| Cave::new(second.into()));
                second_cave.add_connection(first);
            }
        }

        Ok(Self(caves))
    }

    fn part1(&self) -> anyhow::Result<String> {
        let m = HashMap::new();
        Ok(format!("{}", self.search("start", &m, true)))
    }

    fn part2(&self) -> anyhow::Result<String> {
        let m = HashMap::new();
        Ok(format!("{}", self.search("start", &m, false)))
    }
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Copy, Clone)]
enum CaveType {
    Large,
    Small,
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
struct Cave {
    connections: Vec<String>,
    cave_type: CaveType,
}

impl Cave {
    fn add_connection<T>(&mut self, name: T) where T: ToString {
        self.connections.push(name.to_string());
        self.connections.sort();
        self.connections.dedup();
    }

    fn new(cave_name: String) -> Self {
        let cave_type = if cave_name.chars().next().unwrap().is_uppercase() {
            CaveType::Large
        } else {
            CaveType::Small
        };
        Self {
            cave_type,
            connections: vec![],
        }
    }
}

impl Day12 {
    fn search(&self, cur: &str, had: &HashMap<String, i32>, part1: bool) -> i32 {
        let mut had = had.clone();
        if (had.values().any(|&i| i > 1) || part1) && had.contains_key(cur) {
            return 0;
        }
        if cur == "start" && had.contains_key("start") {
            return 0;
        }
        if cur == "end" {
            return 1;
        }
        let n = self.0.get(cur).unwrap();
        if n.cave_type == CaveType::Small {
            let e = had.entry(cur.into()).or_insert(0);
            *e += 1;
        }
        return n.connections.iter().map(|c| self.search(c, &had, part1)).sum::<i32>();
    }
}

#[cfg(test)]
mod tests {
    use crate::{Day, Day12};

    const INPUT_1: &str = r"start-A
start-b
A-c
A-b
b-d
A-end
b-end";

    const INPUT_2: &str = r"dc-end
HN-start
start-kj
dc-start
dc-HN
LN-dc
HN-end
kj-sa
kj-HN
kj-dc";

    const INPUT_3: &str = r"fs-end
he-DX
fs-he
start-DX
pj-DX
end-zg
zg-sl
zg-pj
pj-he
RW-he
fs-DX
pj-RW
zg-RW
start-pj
he-WI
zg-he
pj-fs
start-RW";

    #[test]
    fn part1_test_1() -> anyhow::Result<()> {
        let d = Day12::init(INPUT_1.to_string())?;
        assert_eq!("10", d.part1()?);
        Ok(())
    }

    #[test]
    fn part1_test_2() -> anyhow::Result<()> {
        let d = Day12::init(INPUT_2.to_string())?;
        assert_eq!("19", d.part1()?);
        Ok(())
    }

    #[test]
    fn part1_test_3() -> anyhow::Result<()> {
        let d = Day12::init(INPUT_3.to_string())?;
        assert_eq!("226", d.part1()?);
        Ok(())
    }

    #[test]
    fn part1_real() -> anyhow::Result<()> {
        let d = Day12::init(crate::load_input("12")?)?;
        assert_eq!("4885", d.part1()?);
        Ok(())
    }

    #[test]
    fn part2_test_1() -> anyhow::Result<()> {
        let d = Day12::init(INPUT_1.to_string())?;
        assert_eq!("36", d.part2()?);
        Ok(())
    }

    #[test]
    fn part2_test_2() -> anyhow::Result<()> {
        let d = Day12::init(INPUT_2.to_string())?;
        assert_eq!("103", d.part2()?);
        Ok(())
    }

    // These are very slow, so they are commented out.
    // #[test]
    // fn part2_test_3() -> anyhow::Result<()> {
    //     let d = Day12::init(INPUT_3.to_string())?;
    //     assert_eq!("3509", d.part2()?);
    //     Ok(())
    // }
    //
    // #[test]
    // fn part2_real() -> anyhow::Result<()> {
    //     let d = Day12::init(crate::load_input("12")?)?;
    //     assert_eq!("117095", d.part2()?);
    //     Ok(())
    // }
}