day23: problem 2 solution

23/src/main.rs

@@ -2,7 +2,7 @@ use itertools::Itertools;
 use ndarray::prelude::*;
 use petgraph::algo::all_simple_paths;
 use petgraph::prelude::*;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Debug, Display, Write};
 use std::fs::File;
 use std::io::{BufRead, BufReader, Lines};
@@ -31,25 +31,31 @@ fn main() {
 
 // PARSE
 
-#[derive(Debug, Clone)]
-enum EdgeType {
-    FromPath,
-    FromSlope,
-}
-
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 struct Node {
     c: char,
     pos: Position,
 }
 
+impl Display for Node {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "({},{})", self.pos.0, self.pos.1)
+    }
+}
+
+impl Debug for Node {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "({},{})", self.pos.0, self.pos.1)
+    }
+}
+
 type Position = (usize, usize);
 
 #[derive(Clone)]
 struct ForestMap {
     map: Array2<char>,
     indexes: HashMap<Position, NodeIndex>,
-    graph: DiGraph<Node, EdgeType>,
+    graph: StableDiGraph<Node, u64>,
     start: Position,
     end: Position,
 }
@@ -91,7 +97,7 @@ impl<T: BufRead> From<Lines<T>> for ForestMap {
             map.len_of(Axis(1)) - 1,
         );
 
-        let mut graph = Graph::default();
+        let mut graph = StableGraph::default();
         let mut indexes = HashMap::new();
         for (pos, c) in map.indexed_iter() {
             if *c != '#' {
@@ -99,47 +105,6 @@ impl<T: BufRead> From<Lines<T>> for ForestMap {
             }
         }
 
-        for (pos, c) in map.indexed_iter() {
-            match c {
-                '#' => continue,
-                '.' => {
-                    adjacent_to(&map, pos).iter().for_each(|adj| {
-                        if indexes.contains_key(&adj) {
-                            graph.add_edge(indexes[&pos], indexes[adj], EdgeType::FromPath);
-                        }
-                    });
-                }
-                '^' => {
-                    if let Some(adj) = offset_pos(&map, pos, (0, -1)) {
-                        if indexes.contains_key(&adj) {
-                            graph.add_edge(indexes[&pos], indexes[&adj], EdgeType::FromSlope);
-                        }
-                    }
-                }
-                '>' => {
-                    if let Some(adj) = offset_pos(&map, pos, (1, 0)) {
-                        if indexes.contains_key(&adj) {
-                            graph.add_edge(indexes[&pos], indexes[&adj], EdgeType::FromSlope);
-                        }
-                    }
-                }
-                'v' => {
-                    if let Some(adj) = offset_pos(&map, pos, (0, 1)) {
-                        if indexes.contains_key(&adj) {
-                            graph.add_edge(indexes[&pos], indexes[&adj], EdgeType::FromSlope);
-                        }
-                    }
-                }
-                '<' => {
-                    if let Some(adj) = offset_pos(&map, pos, (-1, 0)) {
-                        if indexes.contains_key(&adj) {
-                            graph.add_edge(indexes[&pos], indexes[&adj], EdgeType::FromSlope);
-                        }
-                    }
-                }
-                c => panic!("invalid map character {}", c),
-            }
-        }
         Self {
             map,
             start,
@@ -150,6 +115,129 @@ impl<T: BufRead> From<Lines<T>> for ForestMap {
     }
 }
 
+impl ForestMap {
+    fn build_graph(&mut self) {
+        for (pos, c) in self.map.indexed_iter() {
+            match c {
+                '#' => continue,
+                '.' => {
+                    adjacent_to(&self.map, pos).iter().for_each(|adj| {
+                        if self.indexes.contains_key(&adj) {
+                            self.graph.add_edge(self.indexes[&pos], self.indexes[adj], 1);
+                        }
+                    });
+                }
+                '^' => {
+                    if let Some(adj) = offset_pos(&self.map, pos, (0, -1)) {
+                        if self.indexes.contains_key(&adj) {
+                            self.graph.add_edge(self.indexes[&pos], self.indexes[&adj], 1);
+                        }
+                    }
+                }
+                '>' => {
+                    if let Some(adj) = offset_pos(&self.map, pos, (1, 0)) {
+                        if self.indexes.contains_key(&adj) {
+                            self.graph.add_edge(self.indexes[&pos], self.indexes[&adj], 1);
+                        }
+                    }
+                }
+                'v' => {
+                    if let Some(adj) = offset_pos(&self.map, pos, (0, 1)) {
+                        if self.indexes.contains_key(&adj) {
+                            self.graph.add_edge(self.indexes[&pos], self.indexes[&adj], 1);
+                        }
+                    }
+                }
+                '<' => {
+                    if let Some(adj) = offset_pos(&self.map, pos, (-1, 0)) {
+                        if self.indexes.contains_key(&adj) {
+                            self.graph.add_edge(self.indexes[&pos], self.indexes[&adj], 1);
+                        }
+                    }
+                }
+                c => panic!("invalid map character {}", c),
+            }
+        }
+    }
+
+    fn build_graph2(&mut self) {
+        for (pos, c) in self.map.indexed_iter() {
+            match c {
+                '#' => continue,
+                '.' | '^' | '>' | 'v' | '<' => {
+                    adjacent_to(&self.map, pos).iter().for_each(|adj| {
+                        if self.indexes.contains_key(&adj) {
+                            self.graph.add_edge(self.indexes[&pos], self.indexes[adj], 1);
+                        }
+                    });
+                }
+                c => panic!("invalid map character {}", c),
+            }
+        }
+    }
+
+    // Cull nodes that don't change the topology of the graph and combine their cost
+    fn simplify_graph(&mut self) {
+        let mut idxs: Vec<_> = self
+            .graph
+            .neighbors(self.indexes[&self.start])
+            .map(|idx| (self.indexes[&self.start], idx))
+            .collect();
+        let mut visited = HashSet::from([self.indexes[&self.start]]);
+        while let Some((last_idx, cur_idx)) = idxs.pop() {
+            if !visited.insert(cur_idx) {
+                continue;
+            }
+            let our_neighbors = self.graph.neighbors(cur_idx).collect_vec();
+
+            // if we have exactly 2 neighbours, then one is where we came from, and we can shortcut this node with a
+            // pair of new edges A <-> C and break the existing 4 edges between them
+            if our_neighbors.len() == 2 {
+                let next_idx = our_neighbors.iter().find(|n| **n != last_idx).unwrap();
+
+                // remove the 4 existing edges
+                // careful of order of operations, as removing edges invalidates edge indexes
+                let forward_cost = self
+                    .graph
+                    .remove_edge(self.graph.find_edge(cur_idx, *next_idx).unwrap())
+                    .unwrap();
+                let last_forward_cost = self
+                    .graph
+                    .remove_edge(self.graph.find_edge(last_idx, cur_idx).unwrap())
+                    .unwrap();
+                let backward_cost = self
+                    .graph
+                    .remove_edge(self.graph.find_edge(cur_idx, last_idx).unwrap())
+                    .unwrap();
+                let next_backward_cost = self
+                    .graph
+                    .remove_edge(self.graph.find_edge(*next_idx, cur_idx).unwrap())
+                    .unwrap();
+                let new_forward_cost = forward_cost + last_forward_cost;
+                let new_backward_cost = backward_cost + next_backward_cost;
+
+                // add edge from last to next
+                self.graph.add_edge(last_idx, *next_idx, new_forward_cost);
+                self.graph.add_edge(*next_idx, last_idx, new_backward_cost);
+
+                self.graph.remove_node(cur_idx);
+                // push the next node
+                idxs.push((last_idx, *next_idx));
+            } else {
+                // don't do anything about nodes with > 2 edges, just push them onto the stack, if there are some
+                idxs.append(
+                    &mut self
+                        .graph
+                        .neighbors(cur_idx)
+                        .into_iter()
+                        .map(|next_idx| (cur_idx, next_idx))
+                        .collect(),
+                );
+            }
+        }
+    }
+}
+
 impl Debug for ForestMap {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         for y in 0..self.map.len_of(Axis(1)) {
@@ -158,20 +246,6 @@ impl Debug for ForestMap {
             }
             writeln!(f)?;
         }
-        // println!("start: {:?} end: {:?}", self.start, self.end);
-        // println!("digraph aoc23 {{");
-        // for node in self.graph.node_indices() {
-        //     println!(
-        //         "  \"{},{}\" -> {}",
-        //         self.graph[node].pos.0,
-        //         self.graph[node].pos.1,
-        //         self.graph
-        //             .neighbors(node)
-        //             .map(|n| format!("\"{},{}\"", self.graph[n].pos.0, self.graph[n].pos.1))
-        //             .join(",")
-        //     );
-        // }
-        // println!("}}");
         Ok(())
     }
 }
@@ -179,8 +253,9 @@ impl Debug for ForestMap {
 // PROBLEM 1 solution
 
 fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
-    let map = ForestMap::from(input);
-    println!("{:?}", map);
+    let mut map = ForestMap::from(input);
+    map.build_graph();
+    // println!("{:?}", map);
     let paths = all_simple_paths::<Vec<_>, _>(&map.graph, map.indexes[&map.start], map.indexes[&map.end], 0, None)
         .collect_vec();
     let longest = paths.iter().max_by_key(|path| path.len()).unwrap();
@@ -188,9 +263,24 @@ fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
     longest.len() as u64 - 1
 }
 
+fn calc_path_length(map: &ForestMap, path: &Vec<NodeIndex>) -> u64 {
+    path.iter().tuple_windows().fold(0, |accum, (prev, next)| {
+        accum + map.graph[map.graph.find_edge(*prev, *next).unwrap()]
+    })
+}
+
 // PROBLEM 2 solution
 fn problem2<T: BufRead>(input: Lines<T>) -> u64 {
-    0
+    let mut map = ForestMap::from(input);
+    map.build_graph2();
+    map.simplify_graph();
+
+    let paths = all_simple_paths::<Vec<_>, _>(&map.graph, map.indexes[&map.start], map.indexes[&map.end], 0, None)
+        .collect_vec();
+    let longest = paths.iter().max_by_key(|path| calc_path_length(&map, &path)).unwrap();
+    longest.iter().tuple_windows().fold(0, |accum, (prev, next)| {
+        accum + map.graph[map.graph.find_edge(*prev, *next).unwrap()]
+    })
 }
 
 #[cfg(test)]
@@ -231,6 +321,6 @@ mod tests {
     #[test]
     fn problem2_example() {
         let c = Cursor::new(EXAMPLE);
-        assert_eq!(problem2(c.lines()), 0);
+        assert_eq!(problem2(c.lines()), 154);
     }
 }