From 49368e79858ddd85df72695fd2c54f878933c6ac Mon Sep 17 00:00:00 2001
From: jazzpi <jasper@mezzo.de>
Date: Wed, 21 Dec 2022 19:54:27 +0100
Subject: [PATCH] Track paths during distance calculation

---
 src/bin/d16p1.rs |  8 ++++----
 src/day16.rs     | 49 +++++++++++++++++++++++++++++++++---------------
 2 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/src/bin/d16p1.rs b/src/bin/d16p1.rs
index 86b2adc..f312b9b 100644
--- a/src/bin/d16p1.rs
+++ b/src/bin/d16p1.rs
@@ -5,13 +5,13 @@ use aoc22::{
 
 pub fn main() {
     let valves = day16::parse_valves(&util::parse_input());
-    let dists = day16::calc_dists(&valves);
-    let state = day16::State::new(&valves);
+    let paths = day16::calc_paths(&valves);
+    let state = day16::State::new(&valves, false);
 
-    let score = util::maximize(&state, &(&valves, &dists));
+    let score = util::maximize(&state, &(&valves, &paths));
 
     println!(
         "Most pressure released is {}",
-        score.lower_bound(&(&valves, &dists))
+        score.lower_bound(&(&valves, &paths))
     );
 }
diff --git a/src/day16.rs b/src/day16.rs
index a1fc649..860275f 100644
--- a/src/day16.rs
+++ b/src/day16.rs
@@ -56,21 +56,40 @@ pub fn parse_valves(input: &String) -> Vec<Valve> {
         .collect()
 }
 
-pub fn calc_dists(valves: &Vec<Valve>) -> Vec<Vec<usize>> {
+#[derive(Debug, Clone)]
+pub struct Path {
+    pub dist: usize,
+    pub path: Vec<usize>,
+}
+
+pub fn calc_paths(valves: &Vec<Valve>) -> Vec<Vec<Path>> {
     let n = valves.len();
-    let mut result = vec![vec![usize::MAX; n]; n];
+
+    // result[from][to] = (dist, path)
+    let mut result = vec![
+        vec![
+            Path {
+                dist: usize::MAX,
+                path: Vec::new()
+            };
+            n
+        ];
+        n
+    ];
 
     for i in 0..n {
-        result[i][i] = 0;
+        result[i][i].dist = 0;
         let mut next = VecDeque::new();
         next.push_back(i);
         while let Some(j) = next.pop_front() {
-            let d_next = result[i][j] + 1;
-            for j in &valves[j].tunnels {
-                let d_tunnel = result[i].get_mut(*j).unwrap();
-                if d_next < *d_tunnel {
-                    *d_tunnel = d_next;
-                    next.push_back(*j);
+            let d_next = result[i][j].dist + 1;
+            for tunnel in &valves[j].tunnels {
+                let d_tunnel = &result[i][*tunnel];
+                if d_next < d_tunnel.dist {
+                    let mut path = result[i][j].path.clone();
+                    path.push(*tunnel);
+                    result[i][*tunnel] = Path { dist: d_next, path };
+                    next.push_back(*tunnel);
                 }
             }
         }
@@ -110,9 +129,9 @@ impl State {
     }
 }
 
-type ExtraArgs<'a> = (&'a Vec<Valve>, &'a Vec<Vec<usize>>);
+type ExtraArgs<'a> = (&'a Vec<Valve>, &'a Vec<Vec<Path>>);
 impl util::BnBState<ExtraArgs<'_>> for State {
-    fn possible_actions(&self, (valves, dists): &ExtraArgs) -> Vec<State> {
+    fn possible_actions(&self, (valves, paths): &ExtraArgs) -> Vec<State> {
         assert!(!self.finished());
 
         let mut result = Vec::new();
@@ -132,7 +151,7 @@ impl util::BnBState<ExtraArgs<'_>> for State {
 
         for l in &self.interesting {
             let i = l.clone();
-            let dist = dists[loc][i];
+            let dist = paths[loc][i].dist;
             if dist == 0 || dist == usize::MAX || self.time_remaining < dist {
                 continue;
             }
@@ -160,21 +179,21 @@ impl util::BnBState<ExtraArgs<'_>> for State {
         self.pressure_released + self.flow_rate * self.time_remaining
     }
 
-    fn upper_bound(&self, (valves, dists): &ExtraArgs) -> usize {
+    fn upper_bound(&self, (valves, paths): &ExtraArgs) -> usize {
         let mut additional_flow = 0;
         for i in &self.interesting {
             let i = i.clone();
             if (self.opened & (1 << i)) != 0 {
                 continue;
             }
-            let dist = dists[self.location][i];
+            let dist = paths[self.location][i].dist;
             if dist < self.time_remaining {
                 let flow = valves[i].flow_rate;
                 additional_flow += flow * (self.time_remaining - dist - 1);
             }
         }
 
-        self.lower_bound(&(valves, dists)) + additional_flow
+        self.lower_bound(&(valves, paths)) + additional_flow
     }
 }