Day 16, part 2

src/bin/d16p1.rs

@@ -3,15 +3,17 @@ use aoc22::{
     util::{self, BnBState},
 };
 
+const TOTAL_TIME: usize = 30;
+
 pub fn main() {
     let valves = day16::parse_valves(&util::parse_input());
     let paths = day16::calc_paths(&valves);
-    let state = day16::State::new(&valves, false);
+    let state = day16::State::new(&valves, TOTAL_TIME, 1);
 
-    let score = util::maximize(&state, &(&valves, &paths));
+    let state = util::maximize(&state, &(&valves, &paths));
 
     println!(
         "Most pressure released is {}",
-        score.lower_bound(&(&valves, &paths))
+        state.lower_bound(&(&valves, &paths))
     );
 }

src/bin/d16p2.rs

@@ -0,0 +1,20 @@
+use aoc22::{
+    day16,
+    util::{self, BnBState},
+};
+
+const TOTAL_TIME: usize = 26;
+const N_AGENTS: usize = 2;
+
+pub fn main() {
+    let valves = day16::parse_valves(&util::parse_input());
+    let paths = day16::calc_paths(&valves);
+    let state = day16::State::new(&valves, TOTAL_TIME, N_AGENTS);
+
+    let state = util::maximize(&state, &(&valves, &paths));
+
+    println!(
+        "Most pressure released is {}",
+        state.lower_bound(&(&valves, &paths))
+    );
+}

src/day16.rs

@@ -1,17 +1,13 @@
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     hash::Hash,
-    sync::Mutex,
 };
 
 use itertools::Itertools;
-use lazy_static::lazy_static;
 use regex::Regex;
 
 use crate::util;
 
-const TOTAL_TIME: usize = 30;
-
 #[derive(Debug, Clone)]
 pub struct Valve {
     pub flow_rate: usize,
@@ -98,20 +94,28 @@ pub fn calc_paths(valves: &Vec<Valve>) -> Vec<Vec<Path>> {
     result
 }
 
+#[derive(Debug, Clone, Hash, PartialEq)]
+pub enum Goal {
+    Idle,
+    NothingReachable,
+    Opening(usize),
+}
+
 #[derive(Debug, Clone)]
 pub struct State {
     pub opened: u64,
-    pub location: usize,
+    pub locations: Vec<usize>,
+    pub goals: Vec<Goal>,
     pub flow_rate: usize,
     pub time_remaining: usize,
     pub pressure_released: usize,
+    pub history: Vec<Vec<usize>>,
 
     interesting: HashSet<usize>,
-    history: Vec<usize>,
 }
 
 impl State {
-    pub fn new(valves: &Vec<Valve>) -> State {
+    pub fn new(valves: &Vec<Valve>, initial_time: usize, agents: usize) -> State {
         let interesting = valves
             .iter()
             .enumerate()
@@ -119,12 +123,146 @@ impl State {
             .collect();
         State {
             opened: 0,
-            location: 0, // "AA" has to be the first valve alphabetically
+            locations: vec![0; agents],
+            goals: vec![Goal::Idle; agents],
             flow_rate: 0,
-            time_remaining: TOTAL_TIME,
+            time_remaining: initial_time,
             pressure_released: 0,
             interesting,
-            history: vec![0],
+            history: vec![vec![0]; agents],
+        }
+    }
+
+    fn find_goals(&self, (_, paths): &ExtraArgs) -> Vec<State> {
+        let target_dists: HashMap<usize, usize> = self
+            .goals
+            .iter()
+            .enumerate()
+            .filter_map(|(i, g)| {
+                if let Goal::Opening(t) = g {
+                    Some((*t, paths[self.locations[i]][*t].dist))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        let new_goals: HashMap<usize, Vec<usize>> = self
+            .goals
+            .iter()
+            .positions(|g| matches!(g, Goal::Idle))
+            .map(|i| {
+                let loc = self.locations[i];
+                (
+                    i,
+                    self.interesting
+                        .iter()
+                        .filter_map(|target| {
+                            let agent_dist = paths[loc][*target].dist;
+                            if agent_dist + 2 >= self.time_remaining {
+                                // Can't reach + open the valve as well as
+                                // letting anything flow through
+                                return None;
+                            }
+
+                            if let Some(dist) = target_dists.get(target) {
+                                if *dist > agent_dist {
+                                    // Someone else is targeting this, but they're slower
+                                    Some(*target)
+                                } else {
+                                    // Someone else is targeting this, and they're faster
+                                    None
+                                }
+                            } else {
+                                // Noone else is targeting this
+                                Some(*target)
+                            }
+                        })
+                        .collect_vec(),
+                )
+            })
+            .collect();
+
+        if new_goals.len() == 1 {
+            let (i, targets) = new_goals.iter().next().unwrap();
+            if targets.is_empty() {
+                let mut state = self.clone();
+                state.goals[*i] = Goal::NothingReachable;
+                return vec![state];
+            }
+
+            targets
+                .iter()
+                .map(|t| {
+                    let mut state = self.clone();
+                    state.goals[*i] = Goal::Opening(*t);
+                    state
+                })
+                .collect()
+        } else if new_goals.len() == 2 {
+            let mut it = new_goals.iter();
+            let (i1, targets1) = it.next().unwrap();
+            let (i2, targets2) = it.next().unwrap();
+            if targets1.len() == 0 {
+                if targets2.len() == 0 {
+                    let mut state = self.clone();
+                    state.goals[*i1] = Goal::NothingReachable;
+                    state.goals[*i2] = Goal::NothingReachable;
+                    return vec![state];
+                }
+
+                return targets2
+                    .iter()
+                    .map(|t| {
+                        let mut state = self.clone();
+                        state.goals[*i1] = Goal::NothingReachable;
+                        state.goals[*i2] = Goal::Opening(*t);
+                        state
+                    })
+                    .collect();
+            } else if targets2.len() == 0 {
+                return targets1
+                    .iter()
+                    .map(|t| {
+                        let mut state = self.clone();
+                        state.goals[*i1] = Goal::Opening(*t);
+                        state.goals[*i2] = Goal::NothingReachable;
+                        state
+                    })
+                    .collect();
+            } else if targets1.len() == 1 && targets1 == targets2 {
+                // Both agents can only reach one and the same goal
+                let mut state = self.clone();
+                let target = targets1[0];
+                if paths[self.locations[*i1]][target].dist < paths[self.locations[*i2]][target].dist
+                {
+                    state.goals[*i1] = Goal::Opening(target);
+                    state.goals[*i2] = Goal::NothingReachable;
+                } else {
+                    state.goals[*i1] = Goal::NothingReachable;
+                    state.goals[*i2] = Goal::Opening(target);
+                }
+                return vec![state];
+            }
+
+            let mut result = Vec::new();
+
+            for target1 in targets1 {
+                for target2 in targets2 {
+                    if target1 == target2 {
+                        continue;
+                    }
+
+                    let mut state = self.clone();
+                    state.goals[*i1] = Goal::Opening(*target1);
+                    state.goals[*i2] = Goal::Opening(*target2);
+                    result.push(state);
+                }
+            }
+
+            result
+        } else {
+            unimplemented!("Can't find goals for >2 agents at the same time");
         }
     }
 }
@@ -134,45 +272,61 @@ impl util::BnBState<ExtraArgs<'_>> for State {
     fn possible_actions(&self, (valves, paths): &ExtraArgs) -> Vec<State> {
         assert!(!self.finished());
 
-        let mut result = Vec::new();
+        let goalless = self
-
+            .goals
-        let loc = self.location;
+            .iter()
-
+            .enumerate()
-        if (self.opened & (1 << loc)) == 0 && self.interesting.contains(&loc) {
+            .filter(|(_, g)| matches!(g, Goal::Idle))
-            let mut open = self.clone();
+            .collect_vec();
-            open.opened |= 1 << loc;
+        if goalless.len() > 0 {
-            // First increase pressure released, then the flow rate
+            return self.find_goals(&(valves, paths));
-            open.pressure_released += open.flow_rate;
-            open.flow_rate += valves[loc].flow_rate;
-            open.time_remaining -= 1;
-            open.interesting.remove(&loc);
-            result.push(open);
         }
 
-        for l in &self.interesting {
+        let min_goal_dist = self
-            let i = l.clone();
+            .goals
-            let dist = paths[loc][i].dist;
+            .iter()
-            if dist == 0 || dist == usize::MAX || self.time_remaining < dist {
+            .enumerate()
-                continue;
+            .filter_map(|(i, g)| {
+                if let Goal::Opening(target) = g {
+                    Some(paths[self.locations[i]][*target].dist)
+                } else {
+                    None
+                }
+            })
+            .min()
+            .unwrap();
+
+        let time_spent = min_goal_dist + 1;
+        let mut completed = self.clone();
+        completed.pressure_released += completed.flow_rate * time_spent;
+        completed.time_remaining -= time_spent;
+        for (i, goal) in self.goals.iter().enumerate() {
+            if let Goal::Opening(goal) = goal {
+                let path = &paths[self.locations[i]][*goal];
+                if path.dist == min_goal_dist {
+                    completed.locations[i] = *goal;
+                    completed.opened |= 1 << goal;
+                    completed.goals[i] = Goal::Idle;
+                    if completed.interesting.remove(goal) {
+                        // Only increase flow rate if no other agent already
+                        // opened this valve
+                        completed.flow_rate += valves[*goal].flow_rate;
+                    }
+                    completed.history[i].push(*goal);
+                } else {
+                    completed.locations[i] = path.path[time_spent - 1];
+                }
             }
-            let mut moved = self.clone();
-            moved.location = i;
-            moved.pressure_released += moved.flow_rate * dist;
-            moved.time_remaining -= dist;
-            moved.history.push(i);
-            result.push(moved);
         }
-
+        return vec![completed];
-        let mut do_nothing = self.clone();
-        do_nothing.pressure_released += do_nothing.flow_rate * self.time_remaining;
-        do_nothing.time_remaining = 0;
-        result.push(do_nothing);
-
-        result
     }
 
     fn finished(&self) -> bool {
-        self.time_remaining == 0 || self.interesting.is_empty()
+        self.time_remaining == 0
+            || self
+                .goals
+                .iter()
+                .all(|g| matches!(g, Goal::NothingReachable)) //self.interesting.is_empty()
     }
 
     fn lower_bound(&self, _: &ExtraArgs) -> usize {
@@ -186,7 +340,12 @@ impl util::BnBState<ExtraArgs<'_>> for State {
             if (self.opened & (1 << i)) != 0 {
                 continue;
             }
-            let dist = paths[self.location][i].dist;
+            let dist = self
+                .locations
+                .iter()
+                .map(|loc| paths[*loc][i].dist)
+                .min()
+                .unwrap();
             if dist < self.time_remaining {
                 let flow = valves[i].flow_rate;
                 additional_flow += flow * (self.time_remaining - dist - 1);
@@ -200,7 +359,8 @@ impl util::BnBState<ExtraArgs<'_>> for State {
 impl Hash for State {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         self.opened.hash(state);
-        self.location.hash(state);
+        self.locations.hash(state);
+        self.goals.hash(state);
         self.time_remaining.hash(state);
         self.pressure_released.hash(state);
     }
@@ -209,7 +369,8 @@ impl Hash for State {
 impl PartialEq for State {
     fn eq(&self, other: &Self) -> bool {
         self.opened == other.opened
-            && self.location == other.location
+            && self.locations == other.locations
+            && self.goals == other.goals
             && self.time_remaining == other.time_remaining
             && self.pressure_released == other.pressure_released
     }