Day11 - part 2

Enfin ! Indice : mémoïsation (merci kmk).

diff --git a/src/day11.rs b/src/day11.rs
index 0304421b30877dcf78b371a162d3996ddd4e0c22..a68832cdf6484b871dcea7e9a484168180456ef4 100644 (file)
--- a/src/day11.rs
+++ b/src/day11.rs
@@ -1,5 +1,6 @@
 use std::error::Error;
 use std::path::Path;
+use std::collections::HashMap;
 
 fn split_in_two(value: u64) -> Option<(u64, u64)> {
     let n_digits = value.ilog10() + 1;
@@ -12,43 +13,45 @@ fn split_in_two(value: u64) -> Option<(u64, u64)> {
     }
 }
 
-fn count_stones(value: u64, blinks: u32) -> u32 {
-    if blinks == 0 {
-        return 1;
-    }
-    else if value == 0 {
-        return count_stones(1, blinks - 1);
-    }
-    else {
-        if let Some((r, l)) = split_in_two(value) {
-            return count_stones(r, blinks - 1) + count_stones(l, blinks - 1);
+fn count_stones(memo: &mut HashMap<(u64, u32), u64>, value: u64, blinks: u32) -> u64 {
+    if !memo.contains_key(&(value, blinks)) {
+        if blinks == 0 {
+            memo.insert((value, blinks), 1);
+        }
+        else if value == 0 {
+            let tmp = count_stones(memo, 1, blinks - 1);
+            memo.insert((value, blinks), tmp);
         }
         else {
-            return count_stones(value * 2024, blinks - 1);
+            if let Some((r, l)) = split_in_two(value) {
+                let tmp = count_stones(memo, r, blinks - 1) + count_stones(memo, l, blinks - 1);
+                memo.insert((value, blinks), tmp);
+            }
+            else {
+                let tmp = count_stones(memo, value * 2024, blinks - 1);
+                memo.insert((value, blinks), tmp);
+            }
         }
     }
+    *memo.get(&(value, blinks)).unwrap()
 }
 
-fn run_part1(input: &str, blinks: u32) -> Result<u32, Box<dyn Error>> {
-    println!("Running {} - part 1", get_day());
-
-    let res = input.split_whitespace()
+fn solve(input: &str, blinks: u32) -> u64 {
+    let mut memo: HashMap<(u64, u32), u64> = HashMap::new();
+    input.split_whitespace()
         .map(|v| v.parse::<u64>().unwrap())
-        .map(|v| count_stones(v, blinks))
-        .sum::<u32>();
+        .map(|v| count_stones(&mut memo, v, blinks))
+        .sum::<u64>()
+}
 
-    Ok(res)
+fn run_part1(input: &str, blinks: u32) -> Result<u64, Box<dyn Error>> {
+    println!("Running {} - part 1", get_day());
+    Ok(solve(input, blinks))
 }
 
-fn run_part2(input: &str, blinks: u32) -> Result<u32, Box<dyn Error>> {
+fn run_part2(input: &str, blinks: u32) -> Result<u64, Box<dyn Error>> {
     println!("Running {} - part 2", get_day());
-
-    let res = input.split_whitespace()
-        .map(|v| v.parse::<u64>().unwrap())
-        .map(|v| count_stones(v, blinks))
-        .sum::<u32>();
-
-    Ok(res)
+    Ok(solve(input, blinks))
 }
 
 pub fn run(input: &str) -> Result<(), Box<dyn Error>> {