refactor(23/2024): extract graph struct

Author: manhunto

src/solutions/year2024/day23.rs

@@ -1,30 +1,19 @@
 use crate::solutions::Solution;
+use crate::utils::graphs::graph::Graph;
 use itertools::Itertools;
-use std::collections::HashMap;
 
 pub struct Day23;
 
 impl Solution for Day23 {
     fn part_one(&self, input: &str) -> String {
-        let mut neighbours: HashMap<&str, Vec<&str>> = HashMap::new();
+        let graph = self.parse(input);
 
-        let connections: Vec<(&str, &str)> = input
-            .lines()
-            .map(|line| {
-                let (a, b) = line.split_once('-').unwrap();
-
-                neighbours.entry(a).or_default().push(b);
-                neighbours.entry(b).or_default().push(a);
-
-                (a, b)
-            })
-            .collect();
-
-        connections
+        graph
+            .edges()
             .iter()
             .flat_map(|(a, b)| {
-                let a_neighbours = neighbours.get(a).unwrap();
-                let b_neighbours = neighbours.get(b).unwrap();
+                let a_neighbours = graph.neighbours(a);
+                let b_neighbours = graph.neighbours(b);
 
                 a_neighbours
                     .iter()
@@ -34,6 +23,7 @@ impl Solution for Day23 {
                         set.sort();
                         set
                     })
+                    .collect::<Vec<[&str; 3]>>()
             })
             .unique()
             .filter(|set| set.iter().any(|c| c.starts_with("t")))
@@ -46,6 +36,20 @@ impl Solution for Day23 {
     }
 }
 
+impl Day23 {
+    fn parse<'a>(&self, input: &'a str) -> Graph<&'a str> {
+        let mut graph: Graph<&str> = Graph::undirected();
+
+        input.lines().for_each(|line| {
+            let (a, b) = line.split_once('-').unwrap();
+
+            graph.add_edge(a, b);
+        });
+
+        graph
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::solutions::year2024::day23::Day23;

src/utils/graphs/graph.rs

@@ -0,0 +1,68 @@
+use std::collections::{HashMap, HashSet};
+use std::hash::Hash;
+
+enum GraphType {
+    #[allow(dead_code)]
+    Directed,
+    Undirected,
+}
+
+pub struct Graph<T> {
+    nodes: HashSet<T>,
+    edges: HashSet<(T, T)>,
+    neighbours: HashMap<T, Vec<T>>,
+    graph_type: GraphType,
+}
+
+impl<T> Graph<T> {
+    pub fn undirected() -> Self {
+        Self {
+            nodes: HashSet::new(),
+            edges: HashSet::new(),
+            neighbours: HashMap::new(),
+            graph_type: GraphType::Undirected,
+        }
+    }
+
+    pub fn add_edge(&mut self, a: T, b: T)
+    where
+        T: Eq + Hash + Copy,
+    {
+        match self.graph_type {
+            GraphType::Directed => self.add_directed_edge(a, b),
+            GraphType::Undirected => self.add_undirected_edge(a, b),
+        }
+    }
+
+    fn add_directed_edge(&mut self, a: T, b: T)
+    where
+        T: Eq + Hash + Copy,
+    {
+        self.nodes.insert(a);
+        self.nodes.insert(b);
+        self.edges.insert((a, b));
+        self.neighbours.entry(a).or_default().push(b);
+    }
+
+    fn add_undirected_edge(&mut self, a: T, b: T)
+    where
+        T: Eq + Hash + Copy,
+    {
+        self.nodes.insert(a);
+        self.nodes.insert(b);
+        self.edges.insert((a, b));
+        self.neighbours.entry(a).or_default().push(b);
+        self.neighbours.entry(b).or_default().push(a);
+    }
+
+    pub fn edges(&self) -> &HashSet<(T, T)> {
+        &self.edges
+    }
+
+    pub fn neighbours(&self, node: &T) -> Vec<T>
+    where
+        T: Eq + Hash + Clone,
+    {
+        self.neighbours.get(node).unwrap_or(&Vec::new()).to_vec()
+    }
+}

src/utils/graphs/mod.rs

@@ -1,5 +1,6 @@
 pub mod a_star;
 pub mod all_paths;
 pub mod dijkstra;
+pub mod graph;
 pub mod longest_path;
 mod state_utils;