Implement day 12 part 2

2021/src/day12.rs

@@ -30,20 +30,30 @@ fn is_small(cave: &str) -> bool {
     cave.chars().all(|c| c.is_ascii_lowercase())
 }
 
-fn dfs_routes<'a>(edges: &'a EdgeMap, route: &'_ mut Vec<&'a str>, pos: &'a str) -> usize {
-    if is_small(pos) && route.contains(&pos) {
+fn dfs_routes<'a>(
+    edges: &'a EdgeMap,
+    route: &'_ mut Vec<&'a str>,
+    pos: &'a str,
+    mut small_twice: bool,
+) -> usize {
+    match pos {
+        "end" => return 1,
+        "start" if !route.is_empty() => return 0,
+        pos if is_small(pos) && route.contains(&pos) => {
+            if small_twice {
+                small_twice = false;
+            } else {
                 return 0;
             }
-
-    if pos == "end" {
-        return 1;
+        }
+        _ => (),
     }
 
     route.push(pos);
 
     let routes = edges[pos]
         .iter()
-        .map(|new_pos| dfs_routes(edges, route, new_pos))
+        .map(|new_pos| dfs_routes(edges, route, new_pos, small_twice))
         .sum();
 
     route.pop();
@@ -51,15 +61,19 @@ fn dfs_routes<'a>(edges: &'a EdgeMap, route: &'_ mut Vec<&'a str>, pos: &'a str)
     routes
 }
 
-pub fn part1(input: &mut dyn Read) -> String {
+fn parts_common(input: &mut dyn Read, small_twice: bool) -> String {
     let edges = read_edges(input);
     let mut route = Vec::new();
 
-    dfs_routes(&edges, &mut route, "start").to_string()
+    dfs_routes(&edges, &mut route, "start", small_twice).to_string()
 }
 
-pub fn part2(_input: &mut dyn Read) -> String {
-    todo!()
+pub fn part1(input: &mut dyn Read) -> String {
+    parts_common(input, false)
+}
+
+pub fn part2(input: &mut dyn Read) -> String {
+    parts_common(input, true)
 }
 
 #[cfg(test)]
@@ -78,4 +92,11 @@ mod tests {
         test_implementation(part1, SAMPLE2, 19);
         test_implementation(part1, SAMPLE3, 226);
     }
+
+    #[test]
+    fn sample_part2() {
+        test_implementation(part2, SAMPLE1, 36);
+        test_implementation(part2, SAMPLE2, 103);
+        test_implementation(part2, SAMPLE3, 3509);
+    }
 }