Non-functional implementation of 2022 day 16 part 1

2022/src/day16.rs

@@ -1,9 +1,246 @@
-use anyhow::Result;
+use std::collections::VecDeque;
 
-pub fn part1(_input: &[u8]) -> Result<String> {
-    anyhow::bail!("not implemented")
+use ahash::AHashMap;
+use ahash::AHashSet;
+use anyhow::Result;
+use nom::branch::alt;
+use nom::bytes::complete::tag;
+use nom::character::complete::alpha1;
+use nom::character::complete::newline;
+use nom::combinator::into;
+use nom::multi::fold_many1;
+use nom::multi::separated_list1;
+use nom::sequence::preceded;
+use nom::sequence::terminated;
+use nom::sequence::tuple;
+use nom::IResult;
+
+use crate::common::parse_input;
+
+type ParsedNetwork<'a> = AHashMap<&'a [u8], ParsedValve<'a>>;
+
+struct ParsedValve<'a> {
+    connected: Vec<&'a [u8]>,
+    flow: u32,
+}
+
+#[derive(Debug)]
+struct SimpleNetwork {
+    valves: Vec<SimpleValve>,
+    start: usize,
+}
+
+#[derive(Debug)]
+struct SimpleValve {
+    connected: Vec<(usize, u8)>,
+    flow: u32,
+}
+
+impl From<ParsedNetwork<'_>> for SimpleNetwork {
+    fn from(parsed: ParsedNetwork) -> Self {
+        let mapping: AHashMap<_, _> = parsed
+            .iter()
+            .filter_map(|(&k, v)| (v.flow > 0 || k == b"AA").then_some(k))
+            .enumerate()
+            .map(|(i, v)| (v, i))
+            .collect();
+
+        let mut todo = VecDeque::new();
+        let mut seen = AHashSet::new();
+
+        let mut network = Vec::with_capacity(mapping.len());
+
+        for (&key, valve_data) in &parsed {
+            if valve_data.flow == 0 && key != b"AA" {
+                continue;
+            }
+
+            todo.extend(valve_data.connected.iter().map(|&valve| (valve, 1)));
+
+            let mut connected = Vec::new();
+
+            seen.clear();
+            while let Some((valve, dist)) = todo.pop_front() {
+                if seen.insert(valve) {
+                    let data = &parsed[&valve];
+
+                    if data.flow != 0 {
+                        connected.push((mapping[valve], dist));
+                    }
+                    for &other in &data.connected {
+                        if other != key {
+                            todo.push_back((other, dist + 1));
+                        }
+                    }
+                }
+            }
+
+            network.push(SimpleValve {
+                flow: valve_data.flow,
+                connected,
+            })
+        }
+
+        Self {
+            valves: network,
+            start: mapping[&b"AA"[..]],
+        }
+    }
+}
+
+fn parse_network(input: &[u8]) -> IResult<&[u8], ParsedNetwork> {
+    let parse_network = terminated(
+        tuple((
+            // Parse the name of the valve
+            preceded(tag("Valve "), alpha1),
+            // Parse the flow of the valve
+            preceded(tag(" has flow rate="), nom::character::complete::u32),
+            // Parse the connections
+            preceded(
+                // Did you really have to distinguish plural
+                alt((
+                    tag("; tunnels lead to valves "),
+                    tag("; tunnel leads to valve "),
+                )),
+                separated_list1(tag(", "), alpha1),
+            ),
+        )),
+        newline,
+    );
+
+    fold_many1(
+        parse_network,
+        ParsedNetwork::new,
+        |mut map, (valve, flow, connected)| {
+            map.insert(valve, ParsedValve { flow, connected });
+
+            map
+        },
+    )(input)
+}
+
+#[derive(Eq, PartialEq, Hash, Debug, Clone)]
+struct State {
+    pos: usize,
+    valves_open: u32,
+}
+
+impl State {
+    fn open(&self) -> Option<State> {
+        let bit = 1 << self.pos;
+        if (self.valves_open & bit) == 0 {
+            Some(State {
+                pos: self.pos,
+                valves_open: self.valves_open | bit,
+            })
+        } else {
+            None
+        }
+    }
+
+    pub fn is_open(&self, pos: usize) -> bool {
+        (self.valves_open & (1 << pos)) != 0
+    }
+}
+
+impl Ord for State {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        // Reversed because having fewer valves with the same score gives more opportunities for gains
+        self.valves_open
+            .count_ones()
+            .cmp(&other.valves_open.count_ones())
+            // Compare open valves and pos. Shouldn't really matter but required for a total order.
+            .then(self.valves_open.cmp(&other.valves_open))
+            .then(self.pos.cmp(&other.pos))
+    }
+}
+
+impl PartialOrd for State {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+pub fn part1(input: &[u8]) -> Result<String> {
+    let network: SimpleNetwork = parse_input(input, into(parse_network))?;
+
+    let mut best = AHashMap::new();
+
+    let initial_state = State {
+        valves_open: 0,
+        pos: network.start,
+    };
+
+    best.insert(initial_state.clone(), 0);
+
+    let mut todo = VecDeque::new();
+
+    todo.push_back((0, 0, initial_state));
+
+    let mut best_score = 0;
+
+    while let Some((score, minute, state)) = todo.pop_front() {
+        if best[&state] > score {
+            continue;
+        }
+
+        let mut enqueue = |score, minute, state: State| {
+            if minute >= 29
+                || best
+                    .get(&state)
+                    .map(|&previous| previous >= score)
+                    .unwrap_or(false)
+            {
+                return;
+            }
+
+            best.insert(state.clone(), score);
+            todo.push_back((score, minute, state));
+        };
+
+        if let Some(new_state) = state.open() {
+            let pos = new_state.pos;
+            let valve_strength = network.valves[pos].flow;
+            let time_remaining = 29 - minute;
+            let new_score = score + time_remaining as u32 * network.valves[pos].flow;
+
+            println!("Opening valve {pos} for {valve_strength} for {time_remaining} minutes: {new_score} ({best_score})");
+
+            best_score = best_score.max(new_score);
+
+            enqueue(new_score, minute + 1, new_state)
+        }
+
+        for &(other, dist) in &network.valves[state.pos].connected {
+            if state.is_open(other) {
+                continue;
+            }
+
+            let new_state = State {
+                pos: other,
+                valves_open: state.valves_open,
+            };
+
+            enqueue(score, minute + dist, new_state);
+        }
+    }
+
+    // Guesses: 1802 (too low)
+    Ok(best_score.to_string())
 }
 
 pub fn part2(_input: &[u8]) -> Result<String> {
     anyhow::bail!("not implemented")
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const SAMPLE: &[u8] = include_bytes!("samples/16.txt");
+
+    #[test]
+    fn sample_part1() {
+        assert_eq!(part1(SAMPLE).unwrap(), "1651");
+    }
+}

2022/src/samples/16.txt

@@ -0,0 +1,10 @@
+Valve AA has flow rate=0; tunnels lead to valves DD, II, BB
+Valve BB has flow rate=13; tunnels lead to valves CC, AA
+Valve CC has flow rate=2; tunnels lead to valves DD, BB
+Valve DD has flow rate=20; tunnels lead to valves CC, AA, EE
+Valve EE has flow rate=3; tunnels lead to valves FF, DD
+Valve FF has flow rate=0; tunnels lead to valves EE, GG
+Valve GG has flow rate=0; tunnels lead to valves FF, HH
+Valve HH has flow rate=22; tunnel leads to valve GG
+Valve II has flow rate=0; tunnels lead to valves AA, JJ
+Valve JJ has flow rate=21; tunnel leads to valve II