More tries at day 16

Cargo.lock

@@ -17,6 +17,7 @@ version = "0.1.0"
 dependencies = [
  "env_logger",
  "itertools",
+ "lazy_static",
  "log",
  "num",
  "regex",
@@ -138,6 +139,12 @@ dependencies = [
  "either",
 ]
 
+[[package]]
+name = "lazy_static"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646"
+
 [[package]]
 name = "libc"
 version = "0.2.138"

Cargo.toml

@@ -11,6 +11,7 @@ itertools = "^0.10"
 log = "^0.4"
 env_logger = "^0.10"
 num = "^0.1"
+lazy_static = "^1.4"
 
 [[bin]]
 name = "d1p1"

src/bin/d16p1.rs

@@ -9,9 +9,9 @@ use std::{
 use aoc22::{day16, util};
 
 pub fn main() {
-    let valves = day16::parse_valves(&util::parse_input());
+    let valves = Arc::new(day16::parse_valves(&util::parse_input()));
-    let dists = day16::calc_dists(&valves);
+    let dists = Arc::new(day16::calc_dists(&valves));
-    let state = day16::State::new(Arc::new(valves), Arc::new(dists));
+    let state = day16::State::new(&valves);
 
     let possible_states = Arc::new(Mutex::new(Vec::new()));
     possible_states.lock().unwrap().push(state);
@@ -23,7 +23,9 @@ pub fn main() {
     for _ in 0..16 {
         let s = possible_states.clone();
         let l = lower_bound.clone();
-        handles.push(thread::spawn(move || check_states(s, l)));
+        let v = valves.clone();
+        let d = dists.clone();
+        handles.push(thread::spawn(move || check_states(s, l, v, d)));
     }
 
     for handle in handles {
@@ -36,8 +38,15 @@ pub fn main() {
     );
 }
 
-pub fn check_states(possible_states: Arc<Mutex<Vec<day16::State>>>, lower_bound: Arc<AtomicUsize>) {
+pub fn check_states(
+    possible_states: Arc<Mutex<Vec<day16::State>>>,
+    lower_bound: Arc<AtomicUsize>,
+    valves: Arc<Vec<day16::Valve>>,
+    dists: Arc<Vec<Vec<usize>>>,
+) {
+    let mut i = 0;
     loop {
+        i += 1;
         let state = { possible_states.lock().unwrap().pop() };
         if state.is_none() {
             break;
@@ -45,14 +54,20 @@ pub fn check_states(possible_states: Arc<Mutex<Vec<day16::State>>>, lower_bound:
         let state = state.unwrap();
         if state.finished() {
             let score = state.lower_bound();
+            dbg!(score);
             lower_bound.fetch_max(score, Ordering::Relaxed);
         } else {
-            let x = state.possible_actions();
+            let x = day16::possible_actions(&state, &valves, &dists);
+            // let x = state.possible_actions(&valves, &dists);
+            // let state_upper = state.upper_bound(&valves, &dists);
             for action in x {
-                if action.upper_bound() > lower_bound.load(Ordering::Relaxed) {
+                let action_upper = action.upper_bound(&valves, &dists);
+                // assert!(action_upper <= state_upper);
+                if action_upper > lower_bound.load(Ordering::Relaxed) {
                     possible_states.lock().unwrap().push(action);
                 }
             }
         }
     }
+    dbg!(i);
 }

src/day16.rs

@@ -1,22 +1,30 @@
 use std::{
-    cmp::Ordering,
     collections::{HashMap, HashSet, VecDeque},
-    sync::Arc,
+    hash::Hash,
+    sync::Mutex,
 };
 
 use itertools::Itertools;
+use lazy_static::lazy_static;
 use regex::Regex;
 
 const TOTAL_TIME: usize = 30;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Valve {
-    pub name: String,
     pub flow_rate: usize,
-    pub tunnels: Vec<String>,
+    pub tunnels: Vec<usize>,
 }
 
-pub fn parse_valves(input: &String) -> HashMap<String, Valve> {
+pub fn parse_valves(input: &String) -> Vec<Valve> {
+    let indices: HashMap<String, usize> = input
+        .lines()
+        .map(|l| l.split(' ').nth(1).unwrap())
+        .sorted()
+        .enumerate()
+        .map(|(i, n)| (n.to_owned(), i))
+        .collect();
+
     let mut valves = HashMap::new();
 
     let re =
@@ -32,41 +40,38 @@ pub fn parse_valves(input: &String) -> HashMap<String, Valve> {
             .unwrap()
             .as_str()
             .split(", ")
-            .map(|s| s.to_owned())
+            .map(|s| indices[s])
             .collect();
-        let valve = Valve {
+        let valve = Valve { flow_rate, tunnels };
-            name: name.clone(),
-            flow_rate,
-            tunnels,
-        };
         valves.insert(name, valve);
     }
 
     valves
+        .iter()
+        .map(|(name, v)| (indices[name], v))
+        .sorted_by_key(|(i, _)| *i)
+        .map(|(_, v)| v.clone())
+        .collect()
 }
 
-pub fn calc_dists(valves: &HashMap<String, Valve>) -> HashMap<(String, String), usize> {
+pub fn calc_dists(valves: &Vec<Valve>) -> Vec<Vec<usize>> {
-    let mut result = HashMap::new();
+    let n = valves.len();
-    let dists_orig: HashMap<String, usize> = valves.keys().map(|name| (name.clone(), 0)).collect();
+    let mut result = vec![vec![usize::MAX; n]; n];
 
-    for name in valves.keys() {
+    for i in 0..n {
-        let mut dists = dists_orig.clone();
+        result[i][i] = 0;
         let mut next = VecDeque::new();
-        next.push_back(name);
+        next.push_back(i);
-        while let Some(n) = next.pop_front() {
+        while let Some(j) = next.pop_front() {
-            let d_next = dists[n] + 1;
+            let d_next = result[i][j] + 1;
-            for tunnel in &valves[n].tunnels {
+            for j in &valves[j].tunnels {
-                let d_tunnel = dists.get_mut(tunnel).unwrap();
+                let d_tunnel = result[i].get_mut(*j).unwrap();
-                if *d_tunnel > d_next {
+                if d_next < *d_tunnel {
                     *d_tunnel = d_next;
-                    next.push_back(tunnel);
+                    next.push_back(*j);
                 }
             }
         }
-
-        for n in valves.keys() {
-            result.insert((name.clone(), n.clone()), dists[n]);
-        }
     }
 
     result
@@ -74,43 +79,29 @@ pub fn calc_dists(valves: &HashMap<String, Valve>) -> HashMap<(String, String),
 
 #[derive(Debug, Clone)]
 pub struct State {
-    pub valves: Arc<HashMap<String, Valve>>,
+    pub opened: u64,
-    pub dists: Arc<HashMap<(String, String), usize>>,
+    pub location: usize,
-    pub opened: HashMap<String, bool>,
-    pub interesting: HashSet<String>,
-    pub location: String,
     pub flow_rate: usize,
     pub time_remaining: usize,
     pub pressure_released: usize,
+
+    interesting: HashSet<usize>,
 }
 
 impl State {
-    pub fn new(
+    pub fn new(valves: &Vec<Valve>) -> State {
-        valves: Arc<HashMap<String, Valve>>,
-        dists: Arc<HashMap<(String, String), usize>>,
-    ) -> State {
-        let opened = valves.iter().map(|(k, _)| (k.clone(), false)).collect();
-
         let interesting = valves
             .iter()
-            .filter_map(|(name, valve)| {
+            .enumerate()
-                if valve.flow_rate > 0 {
+            .filter_map(|(i, v)| if v.flow_rate > 0 { Some(i) } else { None })
-                    Some(name.clone())
-                } else {
-                    None
-                }
-            })
             .collect();
-
         State {
-            valves,
+            opened: 0,
-            dists,
+            location: 0, // "AA" has to be the first valve alphabetically
-            opened,
-            interesting,
-            location: "AA".to_owned(),
             flow_rate: 0,
             time_remaining: TOTAL_TIME,
             pressure_released: 0,
+            interesting,
         }
     }
 
@@ -122,55 +113,51 @@ impl State {
         self.pressure_released + self.flow_rate * self.time_remaining
     }
 
-    pub fn upper_bound(&self) -> usize {
+    pub fn upper_bound(&self, valves: &Vec<Valve>, dists: &Vec<Vec<usize>>) -> usize {
-        let additional_flow: isize = self
+        let mut additional_flow = 0;
-            .opened
+        for i in &self.interesting {
-            .iter()
+            let i = i.clone();
-            .filter_map(|(k, opened)| {
+            if (self.opened & (1 << i)) != 0 {
-                if *opened {
+                continue;
-                    None
+            }
-                } else {
+            let dist = dists[self.location][i];
-                    let dist = self.dists[&(self.location.clone(), k.clone())];
+            if dist < self.time_remaining {
-                    if dist + 1 > self.time_remaining {
+                let flow = valves[i].flow_rate;
-                        None
+                additional_flow += flow * (self.time_remaining - dist - 1);
-                    } else {
+            }
-                        let flow = self.valves[k].flow_rate;
+        }
-                        Some(-((flow * (self.time_remaining - dist - 1)) as isize))
-                    }
-                }
-            })
-            // .sorted()
-            // .rev()
-            // .take(self.time_remaining)
-            .sum();
-        let additional_flow = (-additional_flow) as usize;
 
         self.lower_bound() + additional_flow
     }
 
-    pub fn possible_actions(&self) -> Vec<State> {
+    pub fn possible_actions(&self, valves: &Vec<Valve>, dists: &Vec<Vec<usize>>) -> Vec<State> {
         assert!(!self.finished());
 
         let mut result = Vec::new();
 
-        let loc = &self.location;
+        let loc = self.location;
 
-        if !self.opened[loc] {
+        if (self.opened & (1 << loc)) == 0 && self.interesting.contains(&loc) {
             let mut open = self.clone();
-            *open.opened.get_mut(loc).unwrap() = true;
+            open.opened |= 1 << loc;
             // First increase pressure released, then the flow rate
             open.pressure_released += open.flow_rate;
-            open.flow_rate += self.valves[loc].flow_rate;
+            open.flow_rate += valves[loc].flow_rate;
             open.time_remaining -= 1;
-            open.interesting.remove(loc);
+            open.interesting.remove(&loc);
             result.push(open);
         }
 
-        for tunnel in &self.valves[loc].tunnels {
+        for l in &self.interesting {
+            let i = l.clone();
+            let dist = dists[loc][i];
+            if dist == 0 || dist == usize::MAX || self.time_remaining < dist {
+                continue;
+            }
             let mut moved = self.clone();
-            moved.location = tunnel.clone();
+            moved.location = i;
-            moved.pressure_released += moved.flow_rate;
+            moved.pressure_released += moved.flow_rate * dist;
-            moved.time_remaining -= 1;
+            moved.time_remaining -= dist;
             result.push(moved);
         }
 
@@ -178,22 +165,38 @@ impl State {
     }
 }
 
-impl Ord for State {
+impl Hash for State {
-    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        self.upper_bound().cmp(&other.upper_bound())
+        self.opened.hash(state);
-    }
+        self.location.hash(state);
-}
+        self.time_remaining.hash(state);
-
+        self.pressure_released.hash(state);
-impl PartialOrd for State {
-    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-        Some(self.cmp(other))
     }
 }
 
 impl PartialEq for State {
     fn eq(&self, other: &Self) -> bool {
-        self.cmp(other) == Ordering::Equal
+        self.opened == other.opened
+            && self.location == other.location
+            && self.time_remaining == other.time_remaining
+            && self.pressure_released == other.pressure_released
     }
 }
 
 impl Eq for State {}
+
+pub fn possible_actions(state: &State, valves: &Vec<Valve>, dists: &Vec<Vec<usize>>) -> Vec<State> {
+    lazy_static! {
+        static ref CACHE: Mutex<HashMap<State, Vec<State>>> = Mutex::new(HashMap::new());
+    }
+
+    let actions = { CACHE.lock().unwrap().get(state).cloned() };
+    if let Some(actions) = actions {
+        // assert_eq!(actions, state.possible_actions(valves, dists));
+        actions.clone()
+    } else {
+        let mut cache = CACHE.lock().unwrap();
+        cache.insert(state.clone(), state.possible_actions(valves, dists));
+        cache[state].clone()
+    }
+}