adventofcode22/src/day16.rs

//! Day 16: Proboscidea Volcanium

use std::collections::HashMap;

use once_cell::sync::Lazy;
use regex::Regex;

#[derive(Debug, Clone)]
struct Valve {
    name: String,
    flow: u16,
    connections: Vec<u16>,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct CacheKey {
    open: u64,
    valve: u64,
    remaining: u8,
    pressure: u64,
    do_open: bool,
}

fn parse_input(input: &[String]) -> Vec<Valve> {
    static LINE_REGEX: Lazy<Regex> = Lazy::new(|| {
        Regex::new("Valve (\\w{2}) has flow rate=(\\d+); tunnels? leads? to valves? ([\\w, ]+)")
            .unwrap()
    });

    let ttable = input
        .iter()
        .enumerate()
        .map(|(i, line)| {
            let cap = LINE_REGEX
                .captures(line)
                .unwrap_or_else(|| panic!("could not parse line: {}", line));
            (cap.get(1).unwrap().as_str().to_string(), i as u16)
        })
        .collect::<HashMap<_, _>>();

    input
        .iter()
        .map(|line| {
            let cap = LINE_REGEX
                .captures(line)
                .unwrap_or_else(|| panic!("could not parse line: {}", line));
            Valve {
                name: cap.get(1).unwrap().as_str().to_owned(),
                flow: cap.get(2).unwrap().as_str().parse().unwrap(),
                connections: cap
                    .get(3)
                    .unwrap()
                    .as_str()
                    .split(", ")
                    .map(|n| *ttable.get(n).unwrap())
                    .collect(),
            }
        })
        .collect()
}

fn dfs(
    cache: &mut HashMap<(u64, u16, u16), u16>,
    valves: &[Valve],
    current: usize,
    opened: u64,
    rate: u16,
    steps: u16,
) -> u16 {
    if steps == 0 {
        return 0;
    }

    if let Some(&released) = cache.get(&(opened, steps, current as u16)) {
        return released;
    }

    let mut released = 0;

    // we have 2 options: open the valve or go to any of the connections

    // Can we open the current valve and does it makes sense to do it ?
    if opened & (1 << current) == 0 && valves[current].flow > 0 {
        let next_opened = opened | 1 << current;
        let next_rate = rate + valves[current].flow;
        released = dfs(cache, valves, current, next_opened, next_rate, steps - 1);
    }

    for next in valves[current].connections.iter().copied() {
        let rel = dfs(cache, valves, next as usize, opened, rate, steps - 1);
        released = released.max(rel);
    }

    released += rate;
    cache.insert((opened, steps, current as u16), released);
    released
}

fn dfs_elephant(
    cache: &mut HashMap<(u64, u16, u16, u16), u16>,
    valves: &[Valve],
    me: usize,
    elephant: usize,
    opened: u64,
    rate: u16,
    steps: u16,
) -> u16 {
    if steps == 0 {
        return 0;
    }

    if let Some(&released) = cache.get(&(opened, steps, me as u16, elephant as u16)) {
        return released;
    }

    // both we and the elephant have 2 options: open valve/go next
    let me_can_open = opened & (1 << me) == 0 && valves[me].flow > 0;
    let ele_can_open = me != elephant && opened & (1 << elephant) == 0 && valves[elephant].flow > 0;

    let mut released = 0;

    // I open valve, elephant moves
    if me_can_open {
        let next_opened = opened | 1 << me;
        let next_rate = rate + valves[me].flow;

        for elephant_next in valves[elephant].connections.iter().copied() {
            released = released.max(dfs_elephant(
                cache,
                valves,
                me,
                elephant_next as usize,
                next_opened,
                next_rate,
                steps - 1,
            ));
        }
    }

    // Elephant opens valve, I move
    if ele_can_open {
        let next_opened = opened | 1 << elephant;
        let next_rate = rate + valves[elephant].flow;

        for me_next in valves[me].connections.iter().copied() {
            released = released.max(dfs_elephant(
                cache,
                valves,
                me_next as usize,
                elephant,
                next_opened,
                next_rate,
                steps - 1,
            ));
        }
    }

    // Both open valves
    if me_can_open && ele_can_open {
        let next_opened = opened | 1 << me | 1 << elephant;
        let next_rate = rate + valves[me].flow + valves[elephant].flow;
        released = released.max(dfs_elephant(
            cache,
            valves,
            me,
            elephant,
            next_opened,
            next_rate,
            steps - 1,
        ));
    }

    // Both move
    if !me_can_open || !ele_can_open {
        for me_next in valves[me].connections.iter().copied() {
            for elephant_next in valves[elephant].connections.iter().copied() {
                released = released.max(dfs_elephant(
                    cache,
                    valves,
                    me_next as usize,
                    elephant_next as usize,
                    opened,
                    rate,
                    steps - 1,
                ));
            }
        }
    }

    released += rate;
    cache.insert((opened, steps, me as u16, elephant as u16), released);
    released
}

pub fn part1(input: Vec<String>) -> String {
    let valves = parse_input(&input);
    let mut cache = HashMap::new();

    let valve = valves
        .iter()
        .enumerate()
        .find(|(_, v)| v.name == "AA")
        .unwrap()
        .0;
    let res = dfs(&mut cache, &valves, valve, 0, 0, 30);

    res.to_string()
}

pub fn part2(input: Vec<String>) -> String {
    let valves = parse_input(&input);
    let mut cache = HashMap::new();

    let valve = valves
        .iter()
        .enumerate()
        .find(|(_, v)| v.name == "AA")
        .unwrap()
        .0;
    let res = dfs_elephant(&mut cache, &valves, valve, valve, 0, 0, 26);

    res.to_string()
}

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

    const DAY: u8 = 16;

    #[test]
    fn t_example1() {
        assert_eq!(part1(crate::read_example(DAY)), "1651");
    }

    #[test]
    fn t_example2() {
        assert_eq!(part2(crate::read_example(DAY)), "1707");
    }

    #[test]
    fn t_part1() {
        assert_eq!(part1(crate::read_input(DAY)), "1906");
    }

    #[test]
    fn t_part2() {
        assert_eq!(part2(crate::read_input(DAY)), "2548");
    }
}