[rs] Solve 2022_16 part 2

rs/src/y2022/d16.rs

@@ -54,6 +54,8 @@ impl OpenSet {
     const ALL_CLOSED: Self = Self(0);
 
     /// The powerset of the set of all open valves.
+    ///
+    /// Relies on the fact that the valves are sorted from high to low rate.
     fn powerset(valves: &[Valve]) -> Vec<Self> {
         assert!(valves.len() <= 63);
         let first_zero_rate = valves
@@ -104,16 +106,15 @@ impl Dp1 {
 
 fn solve_part_1(names: &HashMap<&str, usize>, valves: &[Valve], powerset: &[OpenSet]) -> u32 {
     // DP state consists of:
-    // - The current minute
-    // - The current valve
-    // - The current set of open/closed valves
+    // - The valve
+    // - The set of open/closed valves
 
     let mut prev = Dp1::new(valves, powerset);
     let mut curr = Dp1::new(valves, powerset);
 
     // Given valve v in minute t with set s, you can either...
     // - Go to another neighbouring valve v', points: points[v', t-1, s]
-    // - If v not in s: Open current valve, points: points[v, t-1, s+{v}] + t
+    // - If v not in s: Open valve, points: points[v, t-1, s+{v}] + t
 
     for minute in 1..=30 {
         eprintln!("Minute {minute}");
@@ -144,6 +145,112 @@ fn solve_part_1(names: &HashMap<&str, usize>, valves: &[Valve], powerset: &[Open
     curr.get(names["AA"], OpenSet::ALL_CLOSED)
 }
 
+#[derive(Debug)]
+struct Dp2 {
+    valves: usize,
+    powers: usize,
+    elems: Vec<u32>,
+}
+
+impl Dp2 {
+    fn new(valves: &[Valve], powerset: &[OpenSet]) -> Self {
+        let valves = valves.len();
+        let powers = powerset.len();
+        let elems = vec![0; valves * valves * powers];
+        Self {
+            valves,
+            powers,
+            elems,
+        }
+    }
+
+    fn get(&self, own_id: usize, el_id: usize, open: OpenSet) -> u32 {
+        let id = self.powers * (self.valves * own_id + el_id) + open.0 as usize;
+        self.elems[id]
+    }
+
+    fn set(&mut self, own_id: usize, el_id: usize, open: OpenSet, score: u32) {
+        let id = self.powers * (self.valves * own_id + el_id) + open.0 as usize;
+        self.elems[id] = score;
+    }
+
+    fn clear(&mut self) {
+        self.elems.fill(0);
+    }
+}
+
+fn solve_part_2(names: &HashMap<&str, usize>, valves: &[Valve], powerset: &[OpenSet]) -> u32 {
+    // DP state consists of:
+    // - The own valve
+    // - The elephant valve
+    // - The set of open/closed valves
+
+    let mut prev = Dp2::new(valves, powerset);
+    let mut curr = Dp2::new(valves, powerset);
+
+    // Given valve v in minute t with set s, you can either...
+    // - I move, elephant moves
+    // - I move, elephant opens valve
+    // - I open valve, elephant moves
+    // - I open valve, elephant opens valve (if not in same room)
+
+    for minute in 1..=26 {
+        eprintln!("Minute {minute}");
+
+        mem::swap(&mut curr, &mut prev);
+        curr.clear();
+
+        for (own_id, own_valve) in valves.iter().enumerate() {
+            for (el_id, el_valve) in valves.iter().enumerate() {
+                for open in powerset {
+                    let mut score = 0;
+
+                    // Both move
+                    for own_next in &own_valve.next {
+                        for el_next in &el_valve.next {
+                            score = score.max(prev.get(*own_next, *el_next, *open));
+                        }
+                    }
+
+                    // I open valve
+                    let i_can_open = own_valve.rate > 0 && !open.is_open(own_id);
+                    if i_can_open {
+                        let open = open.open(own_id);
+                        for el_next in &el_valve.next {
+                            let room_but_valve_open = prev.get(own_id, *el_next, open);
+                            let pressure_until_end = (minute - 1) * own_valve.rate;
+                            score = score.max(room_but_valve_open + pressure_until_end);
+                        }
+                    }
+
+                    // Elephant opens valve
+                    let el_can_open = el_valve.rate > 0 && !open.is_open(el_id);
+                    if el_can_open {
+                        let open = open.open(el_id);
+                        for own_next in &own_valve.next {
+                            let room_but_valve_open = prev.get(*own_next, el_id, open);
+                            let pressure_until_end = (minute - 1) * el_valve.rate;
+                            score = score.max(room_but_valve_open + pressure_until_end);
+                        }
+                    }
+
+                    // Both open valve
+                    if own_id != el_id && i_can_open && el_can_open {
+                        let open = open.open(own_id).open(el_id);
+                        let room_but_valves_open = prev.get(own_id, el_id, open);
+                        let pressure_until_end = (minute - 1) * (own_valve.rate + el_valve.rate);
+                        score = score.max(room_but_valves_open + pressure_until_end);
+                    }
+
+                    curr.set(own_id, el_id, *open, score);
+                }
+            }
+        }
+    }
+
+    curr.get(names["AA"], names["AA"], OpenSet::ALL_CLOSED)
+}
+
 pub fn solve(input: String) {
     let valves = input.lines().map(parse_valve).collect::<Vec<_>>();
 
@@ -153,4 +260,7 @@ pub fn solve(input: String) {
 
     let part1 = solve_part_1(&names, &valves, &powerset);
     println!("Part 1: {part1}");
+
+    let part2 = solve_part_2(&names, &valves, &powerset);
+    println!("Part 2: {part2}");
 }