day11 solution 2

1 files changed, 38 insertions, 52 deletions

diff --git a/src/bin/day11.rs b/src/bin/day11.rs
index fe8cd2e..a3b8282 100644
--- a/src/bin/day11.rs
+++ b/src/bin/day11.rs
@@ -1,72 +1,58 @@
+use std::collections::HashMap;
+
 static DAY: u8 = 11;
 
 fn main() {
     let input = advent::read_lines(DAY);
-    println!("{DAY}a: {}", blink_stones(&input[0]));
-    println!("{DAY}b: {}", 0);
-}
-
-enum Stone {
-    Number(u64),
-    Split(Box<Stone>,Box<Stone>),
+    println!("{DAY}a: {}", blink_stones(&input[0], 25));
+    println!("{DAY}b: {}", blink_stones(&input[0], 75));
 }
 
-impl Stone {
-
-    fn blink_number(number: u64) -> Stone {
-        /* rule 1 */
-        if number == 0 {
-            return Stone::Number(1);
-        }
+fn next_stones(number: u64) -> Vec<u64> {
+    /* rule 1 */
+    if number == 0 {
+        return [1].to_vec();
+    }
 
-        /* rule 2 */
-        let number_str = number.to_string();
-        if number_str.len() % 2 == 0 {
-            let left = number_str.chars().take(number_str.len() / 2).collect::<String>();
-            let right = number_str.chars().skip(number_str.len() / 2).collect::<String>();
+    /* rule 2 */
+    let number_str = number.to_string();
+    if number_str.len() % 2 == 0 {
+        let left = number_str.chars().take(number_str.len() / 2).collect::<String>();
+        let right = number_str.chars().skip(number_str.len() / 2).collect::<String>();
 
-            let left = left.parse::<u64>().unwrap();
-            let right = right.parse::<u64>().unwrap();
-            return Stone::Split(Box::new(Stone::Number(left)), Box::new(Stone::Number(right)));
-        }
+        let left = left.parse::<u64>().unwrap();
+        let right = right.parse::<u64>().unwrap();
 
-        /* rule 3 */
-        Stone::Number(2024 * number)
+        return [left, right].to_vec();
     }
 
-    fn blink_stone(stone: &Stone) -> Stone {
-        match stone {
-            Stone::Number(number) => Self::blink_number(*number),
-            Stone::Split(left, right) => Stone::Split(Box::new(Self::blink_stone(left)), Box::new(Self::blink_stone(right))),
-        }
+    /* rule 3 */
+    [2024 * number].to_vec()
+}
+
+fn count_stones(number: u64, rounds: u32, cache: &mut HashMap<(u64, u32), u64>) -> u64 {
+    if let Some(n) = cache.get(&(number, rounds)) {
+        return *n;
     }
 
-    fn blink(&mut self) {
-        *self = Self::blink_stone(self);
+    if rounds == 0 {
+        return 1;
     }
 
-    fn count(&self) -> usize {
-        match self {
-            Stone::Number(_) => 1,
-            Stone::Split(left, right) => left.count() + right.count(),
-        }
+    let mut count = 0;
+    for n in next_stones(number) {
+        count += count_stones(n, rounds - 1, cache);
+        cache.insert((number, rounds), count);
     }
+    count
 }
 
-fn blink_stones(input: &str) -> usize {
-    let mut stones = Vec::new();
-    for val in input.split(" ") {
-        let number = val.parse::<u64>().unwrap();
-        stones.push(Stone::Number(number));
-    }
-    for _ in 0..25 {
-        for stone in stones.iter_mut() {
-            stone.blink();
-        }
-    }
-    stones.iter()
-          .map(|stone| stone.count())
-          .sum()
+fn blink_stones(input: &str, rounds: u32) -> u64 {
+    let mut cache = HashMap::new();
+    input.split(" ")
+         .map(|val| val.parse::<u64>().unwrap())
+         .map(|val| count_stones(val, rounds, &mut cache))
+         .sum()
 }
 
 #[cfg(test)]
@@ -76,6 +62,6 @@ mod tests {
     #[test]
     fn test() {
         let input = "125 17";
-        assert_eq!(blink_stones(input), 55312);
+        assert_eq!(blink_stones(input, 25), 55312);
     }
 }