day 19 part two

Author: mark-meyer

2024/day_19/src/lib.rs

@@ -1,4 +1,3 @@
-
 #[derive(Debug)]
 pub struct Trie {
     children: [Option<Box<Trie>>; 5],
@@ -14,15 +13,14 @@ impl Trie {
     }
 
     fn index_stripe(stripe: char) -> usize {
-        // allos us to use a super sparse children array
+        // allows us to use a super sparse children array
         match stripe {
             'b' => 0,
             'g' => 1,
             'r' => 2,
             'u' => 3,
             'w' => 4,
-            _ => panic!("That towel stripe is not allowed in this onsen")
-
+            _ => panic!("That towel stripe is not allowed in this onsen"),
         }
     }
 
@@ -34,37 +32,36 @@ impl Trie {
             current = current.children[i].get_or_insert_with(|| Box::new(Trie::new()));
         }
         current.is_word = true;
-
     }
 
-    pub fn get_prefixes<'a>(&self, towel: &'a str) -> Vec<&'a str>{
+    pub fn get_prefixes<'a>(&self, towel: &'a str) -> Vec<&'a str> {
         let mut current = self;
-        let mut prefixes:Vec<&str> = vec![];
+        let mut prefixes: Vec<&str> = vec![];
 
-        for (len, stripe) in towel.chars().enumerate(){
+        for (len, stripe) in towel.chars().enumerate() {
             let i = Trie::index_stripe(stripe);
-            
+
             if let Some(next) = &current.children[i] {
                 if next.is_word {
                     prefixes.push(&towel[..=len])
                 }
                 current = next
             } else {
-                return prefixes
+                return prefixes;
             }
         }
         prefixes
     }
 
     pub fn is_possible(&self, towel: &str) -> bool {
         let n = towel.len();
-        let mut memo = vec![false;n+1];
+        let mut memo = vec![false; n + 1];
         memo[0] = true;
         for start in 0..n {
             if memo[start] {
                 for prefix in self.get_prefixes(&towel[start..]) {
                     let end = start + prefix.len();
-                    memo[end] =true;
+                    memo[end] = true;
                     if n == end {
                         return true;
                     }
@@ -73,6 +70,25 @@ impl Trie {
         }
         memo[n]
     }
+
+    pub fn count_possible(&self, towel: &str) -> u64 {
+        // this is basically identical to is_possible
+        // but without the early exit
+        let n = towel.len();
+        let mut memo = vec![0; n + 1];
+        memo[0] = 1;
+        for start in 0..n {
+            if memo[start] > 0 {
+                for prefix in self.get_prefixes(&towel[start..]) {
+                    let end = start + prefix.len();
+                    if end <= n {
+                        memo[end] += memo[start];
+                    }
+                }
+            }
+        }
+        memo[n]
+    }
 }
 
 #[cfg(test)]
@@ -83,23 +99,41 @@ mod test {
         let mut t = Trie::new();
         t.insert("bwg");
         t.insert("bwgr");
-        
+
         assert_eq!(t.get_prefixes("bwgrra"), vec!["bwg", "bwgr"]);
     }
 
     #[test]
     fn test_possible() {
         let mut t = Trie::new();
         ["r", "wr", "b", "g", "bwu", "rb", "gb", "br"]
-        .iter().for_each(|p| t.insert(p));
+            .iter()
+            .for_each(|p| t.insert(p));
 
-        // assert_eq!(t.is_possible("brwrr"), true);
-        // assert_eq!(t.is_possible("bggr"), true);
-        // assert_eq!(t.is_possible("gbbr"), true);
-        // assert_eq!(t.is_possible("rrbgbr"), true);
-       // assert_eq!(t.is_possible("ubwu"), false);
-        // assert_eq!(t.is_possible("bwurrg"), true);
-        // assert_eq!(t.is_possible("brgr"), true);
+        assert_eq!(t.is_possible("brwrr"), true);
+        assert_eq!(t.is_possible("bggr"), true);
+        assert_eq!(t.is_possible("gbbr"), true);
+        assert_eq!(t.is_possible("rrbgbr"), true);
+        assert_eq!(t.is_possible("ubwu"), false);
+        assert_eq!(t.is_possible("bwurrg"), true);
+        assert_eq!(t.is_possible("brgr"), true);
         assert_eq!(t.is_possible("bbrgwb"), false);
     }
-}
+
+    #[test]
+    fn count_possible() {
+        let mut t = Trie::new();
+        ["r", "wr", "b", "g", "bwu", "rb", "gb", "br"]
+            .iter()
+            .for_each(|p| t.insert(p));
+
+        assert_eq!(t.count_possible("brwrr"), 2);
+        assert_eq!(t.count_possible("bggr"), 1);
+        assert_eq!(t.count_possible("gbbr"), 4);
+        assert_eq!(t.count_possible("rrbgbr"), 6);
+        assert_eq!(t.count_possible("ubwu"), 0);
+        assert_eq!(t.count_possible("bwurrg"), 1);
+        assert_eq!(t.count_possible("brgr"), 2);
+        assert_eq!(t.count_possible("bbrgwb"), 0);
+    }
+}

2024/day_19/src/main.rs

@@ -1,30 +1,27 @@
-use std::fs;
 use day_19::*;
+use std::fs;
 
-fn parse(s:&str) -> (Trie, Vec<String>) {
+fn parse(s: &str) -> (Trie, Vec<String>) {
     let (patterns, designs) = s.split_once("\n\n").unwrap();
     let mut prefixes = Trie::new();
     for p in patterns.split(", ") {
         prefixes.insert(p.trim());
-    };
-    let designs =  designs.split("\n")
-    .map(|line| line.to_owned())
-    .collect();
+    }
+    let designs = designs.split("\n").map(|line| line.to_owned()).collect();
 
     (prefixes, designs)
 }
 
-fn part_one<'a>(prefixes: &Trie, towels: &[String]) -> usize {
-    towels
-    .iter()
-    .filter(|t| prefixes.is_possible(t))
-    .count()
-
+fn part_one(prefixes: &Trie, towels: &[String]) -> usize {
+    towels.iter().filter(|t| prefixes.is_possible(t)).count()
+}
+fn part_two(prefixes: &Trie, towels: &[String]) -> u64 {
+    towels.iter().map(|t| prefixes.count_possible(t)).sum()
 }
 
 fn main() {
     let s = fs::read_to_string("data.txt").expect("we lost the designs file!");
     let (prefixes, designs) = parse(&s);
     println!("{:?}", part_one(&prefixes, &designs));
-    println!("{:?}", prefixes.is_possible("gwwgwbgbgbuugwurgggwrubrruuwgbwgwrgwrbwrugwwrrugrwgu"));
+    println!("{:?}", part_two(&prefixes, &designs));
 }