Day 14 part 2

1 files changed, 55 insertions, 30 deletions

diff --git a/src/bin/day_14.rs b/src/bin/day_14.rs
index 4814855..e757faf 100644
--- a/src/bin/day_14.rs
+++ b/src/bin/day_14.rs
@@ -11,16 +11,12 @@ use std::{collections::BTreeMap, fs};
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let input = fs::read_to_string("inputs/day_14.txt")?;
     let mut polymer = parse_polymer_expansion(&input).unwrap().1;
-    for _ in 0..10 {
-        polymer.expand();
-        dbg!(&polymer.element_counts);
-    }
-    dbg!(polymer.most_common_element_count() - polymer.least_common_element_count());
-    for _ in 10..40 {
-        polymer.expand();
-        dbg!(&polymer.element_counts);
-    }
-    dbg!(polymer.most_common_element_count() - polymer.least_common_element_count());
+
+    let small_expansion_counts = polymer.count_after_expansions(10);
+    dbg!(small_expansion_counts.most_common() - small_expansion_counts.least_common());
+
+    let big_expansion_counts = polymer.count_after_expansions(40);
+    dbg!(big_expansion_counts.most_common() - big_expansion_counts.least_common());
 
     Ok(())
 }
@@ -29,42 +25,71 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 struct PolymerExpansion {
     polymer: Vec<char>,
     rules: Rules,
-    element_counts: BTreeMap<char, u64>,
+    cache: BTreeMap<(char, char, usize), ElementCount>,
 }
 
 impl PolymerExpansion {
     fn new(polymer: Vec<char>, rules: Rules) -> PolymerExpansion {
-        let mut element_counts = BTreeMap::new();
-        for c in &polymer {
-            *element_counts.entry(*c).or_insert(0) += 1;
-        }
-
         PolymerExpansion {
             polymer,
             rules,
-            element_counts,
+            cache: BTreeMap::new(),
         }
     }
 
-    fn expand(&mut self) {
-        let mut next_polymer = Vec::new();
+    fn count_after_expansions(&mut self, expansions: usize) -> ElementCount {
+        let mut counts = ElementCount::default();
         for i in 0..self.polymer.len() - 1 {
-            next_polymer.push(self.polymer[i]);
-            if let Some(extra) = self.rules.get(&[self.polymer[i], self.polymer[i + 1]]) {
-                next_polymer.push(extra);
-                *self.element_counts.entry(extra).or_insert(0) += 1;
-            }
+            counts.add(self.polymer[i]);
+            counts.append(self.count_between(self.polymer[i], self.polymer[i + 1], expansions));
+        }
+        counts.add(self.polymer[self.polymer.len() - 1]);
+        counts
+    }
+
+    fn count_between(
+        &mut self,
+        left: char,
+        right: char,
+        remaining_expansions: usize,
+    ) -> ElementCount {
+        if remaining_expansions == 0 {
+            ElementCount::default()
+        } else if let Some(cached) = self.cache.get(&(left, right, remaining_expansions)) {
+            cached.clone()
+        } else {
+            let mut counts = ElementCount::default();
+            let middle = self.rules.get(&[left, right]).expect("No rule!");
+            counts.add(middle);
+            counts.append(self.count_between(left, middle, remaining_expansions - 1));
+            counts.append(self.count_between(middle, right, remaining_expansions - 1));
+            self.cache
+                .insert((left, right, remaining_expansions), counts.clone());
+            counts
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone)]
+struct ElementCount(BTreeMap<char, u64>);
+
+impl ElementCount {
+    fn add(&mut self, c: char) {
+        *self.0.entry(c).or_insert(0) += 1;
+    }
+
+    fn append(&mut self, other: ElementCount) {
+        for (key, val) in other.0 {
+            *self.0.entry(key).or_insert(0) += val;
         }
-        next_polymer.push(self.polymer[self.polymer.len() - 1]);
-        self.polymer = next_polymer;
     }
 
-    fn most_common_element_count(&self) -> u64 {
-        self.element_counts.values().max().cloned().unwrap_or(0)
+    fn most_common(&self) -> u64 {
+        self.0.values().max().cloned().unwrap_or(0)
     }
 
-    fn least_common_element_count(&self) -> u64 {
-        self.element_counts.values().min().cloned().unwrap_or(0)
+    fn least_common(&self) -> u64 {
+        self.0.values().min().cloned().unwrap_or(0)
     }
 }