Complete day 11 part 2 and tidy up

diff --git a/src/day11.rs b/src/day11.rs
index 52a88cef5aa45abbdeb0f0dec8f99f391226d591..53346748abbc3cffa47b82d5406fa6b257104aca 100644 (file)
--- a/src/day11.rs
+++ b/src/day11.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 fn parse(input: &str) -> Vec<u64> {
     input
         .split_whitespace()
@@ -5,44 +7,47 @@ fn parse(input: &str) -> Vec<u64> {
         .collect()
 }
 
-fn iterate(stones: &[u64]) -> Vec<u64> {
-    stones
-        .iter()
-        .flat_map(|n| {
-            if *n == 0 {
-                vec![1]
-            } else {
-                let s = format!("{}", n);
-                let l = s.len();
-                if l % 2 == 0 {
-                    vec![s[0..l / 2].parse().unwrap(), s[l / 2..l].parse().unwrap()]
-                } else {
-                    vec![n * 2024]
-                }
-            }
-        })
-        .collect()
+fn succ(stone: u64) -> Vec<u64> {
+    if stone == 0 {
+        vec![1]
+    } else {
+        let s = format!("{}", stone);
+        let l = s.len();
+        if l % 2 == 0 {
+            vec![s[0..l / 2].parse().unwrap(), s[l / 2..l].parse().unwrap()]
+        } else {
+            vec![stone * 2024]
+        }
+    }
+}
+
+fn itercount(stones: &[u64], n: u64, cache: &mut HashMap<(u64, u64), u64>) -> u64 {
+    if n == 0 {
+        stones.len().try_into().unwrap()
+    } else if stones.len() == 1 {
+        let s = stones[0];
+        if let Some(&ct) = cache.get(&(s, n - 1)) {
+            ct
+        } else {
+            let ct = itercount(&succ(s), n - 1, cache);
+            cache.insert((s, n - 1), ct);
+            ct
+        }
+    } else {
+        stones
+            .iter()
+            .copied()
+            .map(|s| itercount(&succ(s), n - 1, cache))
+            .sum()
+    }
 }
 
 pub fn part1(input: &str) -> u64 {
-    std::iter::successors(Some(parse(input)), |v| Some(iterate(v)))
-        .take(26)
-        .last()
-        .unwrap()
-        .len()
-        .try_into()
-        .unwrap()
+    itercount(&parse(input), 25, &mut HashMap::new())
 }
 
 pub fn part2(input: &str) -> u64 {
-    // lol no
-    std::iter::successors(Some(parse(input)), |v| Some(iterate(v)))
-        .take(76)
-        .last()
-        .unwrap()
-        .len()
-        .try_into()
-        .unwrap()
+    itercount(&parse(input), 75, &mut HashMap::new())
 }
 
 #[cfg(test)]
@@ -50,6 +55,7 @@ mod test {
     use super::*;
 
     const INPUT_STR: &str = concat!("0 1 10 99 999\n");
+    const TINY_INPUT: &str = concat!("125 17\n");
 
     #[test]
     fn test_parse() {
@@ -57,12 +63,21 @@ mod test {
     }
 
     #[test]
-    fn test_iterate() {
-        assert_eq!(iterate(&parse(INPUT_STR)), [1, 2024, 1, 0, 9, 9, 2021976]);
+    fn test_succ() {
         assert_eq!(
-            std::iter::successors(Some(parse("125 17")), |v| Some(iterate(v)))
-                .take(7)
+            parse(INPUT_STR)
+                .iter()
+                .copied()
+                .flat_map(succ)
                 .collect::<Vec<_>>(),
+            [1, 2024, 1, 0, 9, 9, 2021976]
+        );
+        assert_eq!(
+            std::iter::successors(Some(parse(TINY_INPUT)), |v| Some(
+                v.iter().copied().flat_map(succ).collect()
+            ))
+            .take(7)
+            .collect::<Vec<_>>(),
             [
                 vec![125, 17],
                 vec![253000, 1, 7],
@@ -78,14 +93,21 @@ mod test {
         )
     }
 
+    #[test]
+    fn test_itercount() {
+        assert_eq!(
+            itercount(&parse(TINY_INPUT), 25, &mut HashMap::new()),
+            55312
+        )
+    }
+
     #[test]
     fn test_part1() {
         assert_eq!(part1(&crate::input(11).unwrap()), 212655)
     }
 
     #[test]
-    #[ignore]
     fn test_part2() {
-        assert_eq!(part2(&crate::input(11).unwrap()), 0)
+        assert_eq!(part2(&crate::input(11).unwrap()), 253582809724830)
     }
 }