[2023] Day 23: Complete part b

aoc_2023/src/day_23.rs

@@ -1,4 +1,7 @@
-use std::{collections::HashSet, convert::identity};
+use std::{
+    collections::{HashMap, HashSet},
+    convert::identity,
+};
 
 use common::{Answer, Solution};
 use nd_vec::{vector, Vec2};
@@ -15,35 +18,15 @@ impl Solution for Day23 {
     }
 
     fn part_a(&self, input: &str) -> Answer {
-        longest_path(
-            &parse(input),
-            Box::new(HashSet::new()),
-            true,
-            vector!(1, 0),
-            0,
-        )
-        .into()
+        solve_a(&parse(input), &mut HashSet::new(), vector!(1, 0), 0).into()
     }
 
     fn part_b(&self, input: &str) -> Answer {
-        longest_path(
-            &parse(input),
-            Box::new(HashSet::new()),
-            false,
-            vector!(1, 0),
-            0,
-        )
-        .into()
+        solve_b(&parse(input)).into()
     }
 }
 
-fn longest_path(
-    map: &Matrix<char>,
-    visited: Box<HashSet<Pos>>,
-    respect_slopes: bool,
-    pos: Pos,
-    idx: u32,
-) -> u32 {
+fn solve_a(map: &Matrix<char>, visited: &mut HashSet<Pos>, pos: Pos, idx: u32) -> u32 {
     if pos == map.size() - vector!(2, 1) {
         return idx;
     }
@@ -59,31 +42,118 @@ fn longest_path(
         }
 
         let next = map[new_pos];
-        if !(if respect_slopes {
-            match dir {
-                Direction::Up => next == '^',
-                Direction::Down => next == 'v',
-                Direction::Left => next == '<',
-                Direction::Right => next == '>',
-            }
-        } else {
-            next == '^' || next == 'v' || next == '<' || next == '>'
+        if !(match dir {
+            Direction::Up => next == '^',
+            Direction::Down => next == 'v',
+            Direction::Left => next == '<',
+            Direction::Right => next == '>',
         } || next == '.')
         {
             continue;
         }
 
-        let mut visited = visited.clone();
         if !visited.insert(pos) {
             continue;
         }
 
-        longest = longest.max(longest_path(map, visited, respect_slopes, new_pos, idx + 1));
+        longest = longest.max(solve_a(map, visited, new_pos, idx + 1));
+        visited.remove(&pos);
     }
 
     longest
 }
 
+fn solve_b(map: &Matrix<char>) -> u32 {
+    // Build graph
+    let mut graph = HashMap::new();
+    for y in 0..map.size().y() {
+        for x in 0..map.size().x() {
+            let pos = vector!(x, y);
+            let c = map[pos];
+            if !b".>v".contains(&(c as u8)) {
+                continue;
+            }
+
+            for dir in Direction::ALL {
+                if let Some(new_pos) = dir.try_advance(pos) {
+                    if new_pos.x() < map.size.x()
+                        && new_pos.y() < map.size.y()
+                        && b".>v".contains(&(map[new_pos] as u8))
+                    {
+                        graph
+                            .entry(pos)
+                            .or_insert_with(HashSet::new)
+                            .insert((new_pos, 1));
+                        graph
+                            .entry(new_pos)
+                            .or_insert_with(HashSet::new)
+                            .insert((pos, 1));
+                    }
+                }
+            }
+        }
+    }
+
+    // Collapse graph
+    let mut dirty = true;
+    while dirty {
+        dirty = false;
+        for key in graph.keys().copied().collect::<Vec<_>>() {
+            if graph[&key].len() != 2 {
+                continue;
+            }
+
+            let ((a, a_score), (b, b_score)) = {
+                let mut iter = graph[&key].iter();
+                let a = iter.next().unwrap();
+                let b = iter.next().unwrap();
+                (*a, *b)
+            };
+
+            let a_graph = graph.get_mut(&a).unwrap();
+            a_graph.retain(|(pos, _)| *pos != key);
+            a_graph.insert((b, a_score + b_score));
+
+            let b_graph = graph.get_mut(&b).unwrap();
+            b_graph.retain(|(pos, _)| *pos != key);
+            b_graph.insert((a, a_score + b_score));
+
+            graph.remove(&key);
+            dirty = true;
+        }
+    }
+
+    // Find longest path
+
+    let mut queue = Vec::<(Pos, Option<u32>)>::new();
+    let mut visited = HashSet::new();
+    let mut max = 0;
+
+    queue.push((vector!(1, 0), Some(0)));
+    while let Some((pos, distance)) = queue.pop() {
+        let Some(distance) = distance else {
+            visited.remove(&pos);
+            continue;
+        };
+
+        if pos == map.size() - vector!(2, 1) {
+            max = max.max(distance);
+            continue;
+        }
+
+        if !visited.insert(pos) {
+            continue;
+        }
+
+        queue.push((pos, None));
+        for (pos, dist) in &graph[&pos] {
+            queue.push((*pos, Some(distance + dist)));
+        }
+    }
+
+    max
+}
+
 fn parse(input: &str) -> Matrix<char> {
     Matrix::new_chars(input, identity)
 }