Day 12, part 2, done at last

1 files changed, 134 insertions, 304 deletions

diff --git a/2023/src/bin/day_12.rs b/2023/src/bin/day_12.rs
index 7485aae..e29042a 100644
--- a/2023/src/bin/day_12.rs
+++ b/2023/src/bin/day_12.rs
@@ -6,7 +6,7 @@ use nom::{
     sequence::separated_pair,
     IResult,
 };
-use std::fs;
+use std::{collections::BTreeMap, fs};
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let input = fs::read_to_string("inputs/day_12.txt")?;
@@ -22,7 +22,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 #[derive(Debug)]
 struct SpringField(Vec<SpringRow>);
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 struct SpringRow {
     springs: Vec<Spring>,
     check: Vec<u32>,
@@ -47,7 +47,7 @@ impl SpringField {
     }
 
     fn possibilities_sum(&self) -> usize {
-        self.0.iter().map(|r| r.possibilities_count_alt()).sum()
+        self.0.iter().map(|r| r.possibilities_count()).sum()
     }
 
     fn expand(&self) -> SpringField {
@@ -85,219 +85,95 @@ impl SpringRow {
         expanded
     }
 
-    /// fail if it's impossible with that prefix to find any matches. Result is (valid, failed-after-reading-x)
-    fn validate_check_prefix(&self, bad_springs: &[usize]) -> (bool, usize) {
-        let mut check_iter = self.check.iter();
-        let mut current_check = check_iter.next();
-        let mut bad_spring_iter = bad_springs.iter();
-        let mut current_bad_spring = bad_spring_iter.next();
-        let mut bad_spring_read_count = 0;
-        let mut current_count = 0;
-        let mut in_bad_lands = false;
-
-        for (i, spring) in self.springs.iter().enumerate() {
-            let (bad, increment) = match spring {
-                Spring::Good => (false, 1),
-                Spring::Bad(s) => (true, *s),
-                Spring::Unknown => {
-                    if current_bad_spring.is_none() {
-                        // That's the end of the prefix. Time to exit early.
-                        let too_many_bads = in_bad_lands
-                            && current_check
-                                .map_or(false, |expected| *expected < current_count as u32);
-                        return (!too_many_bads, bad_spring_read_count);
-                    }
+    fn possibilities_count(&self) -> usize {
+        let optimized = self.optimize_tail();
+        dbg!(optimized.possibilities_count_inner(0, 0, &mut BTreeMap::new()))
+    }
 
-                    let bad_check = current_bad_spring == Some(&i);
-                    if bad_check {
-                        current_bad_spring = bad_spring_iter.next();
-                        bad_spring_read_count += 1;
-                    }
-                    (bad_check, 1)
+    fn optimize_tail(&self) -> SpringRow {
+        let mut optimized = self.clone();
+        while optimized.springs.len() > 0 && optimized.check.len() > 0 {
+            match optimized.springs[optimized.springs.len() - 1] {
+                Spring::Good => {
+                    optimized.springs.pop();
                 }
-            };
-
-            if bad {
-                current_count += increment;
-                in_bad_lands = true;
-            } else {
-                if in_bad_lands {
-                    let valid =
-                        current_check.map_or(false, |expected| *expected == current_count as u32);
-                    if !valid {
-                        return (false, bad_spring_read_count);
-                    }
-
-                    current_count = 0;
-                    current_check = check_iter.next();
-                    in_bad_lands = false;
+                Spring::Bad(s) if s > optimized.check[optimized.check.len() - 1] as usize => {
+                    panic!("Unsolvable row");
+                }
+                Spring::Bad(s) if s <= optimized.check[optimized.check.len() - 1] as usize => {
+                    optimized.trim_bad_suffix();
+                }
+                Spring::Bad(_) => unreachable!(),
+                Spring::Unknown => {
+                    break;
                 }
             }
         }
 
-        (true, bad_spring_read_count)
+        optimized
     }
 
-    fn validate_check(&self, bad_springs: &[usize]) -> bool {
-        let mut check_iter = self.check.iter();
-        let mut current_check = check_iter.next();
-        let mut bad_spring_iter = bad_springs.iter();
-        let mut current_bad_spring = bad_spring_iter.next();
-        let mut current_count = 0;
-        let mut in_bad_lands = false;
-
-        for (i, spring) in self.springs.iter().enumerate() {
-            let (bad, increment) = match spring {
-                Spring::Good => (false, 1),
-                Spring::Bad(s) => (true, *s),
+    fn trim_bad_suffix(&mut self) {
+        let last_check = self.check.pop().expect("No trailing check to pop") as usize;
+        let mut check_i = 0;
+        while check_i < last_check {
+            match self
+                .springs
+                .pop()
+                .expect("Ran out of springs while optimizing suffix!")
+            {
                 Spring::Unknown => {
-                    let bad_check = current_bad_spring == Some(&i);
-                    if bad_check {
-                        current_bad_spring = bad_spring_iter.next();
-                    }
-                    (bad_check, 1)
+                    check_i += 1;
                 }
-            };
-
-            if bad {
-                current_count += increment;
-                in_bad_lands = true;
-            } else {
-                if in_bad_lands {
-                    let valid =
-                        current_check.map_or(false, |expected| *expected == current_count as u32);
-                    if !valid {
-                        return false;
-                    }
-                    current_count = 0;
-                    current_check = check_iter.next();
-                    in_bad_lands = false;
+                Spring::Bad(inc) => {
+                    check_i += inc;
+                }
+                Spring::Good => {
+                    panic!("Found a good spring in the middle of what must be a bad spring range");
                 }
             }
         }
+        let is_boundary = self.springs.len() == 0;
+        let has_extra_good_or_unknown_to_trim =
+            !is_boundary && !matches!(self.springs.last(), Some(Spring::Bad(_)));
 
-        if in_bad_lands {
-            let valid = current_check.map_or(false, |expected| *expected == current_count as u32);
-            if !valid {
-                return false;
-            }
-            current_check = check_iter.next();
+        let valid_suffix =
+            check_i == last_check && (is_boundary || has_extra_good_or_unknown_to_trim);
+        if !valid_suffix {
+            panic!("The suffix had another invalid bad section immediately!");
         }
 
-        current_check == None
-    }
-
-    fn possibilities_count(&self) -> usize {
-        let unknown_indexes: Vec<usize> = self
-            .springs
-            .iter()
-            .enumerate()
-            .filter_map(|(i, s)| (s == &Spring::Unknown).then_some(i))
-            .collect();
-
-        let total_bad_count: usize = self.check.iter().sum::<u32>() as usize;
-        let known_bad_count: usize = self
-            .springs
-            .iter()
-            .map(|s| match s {
-                Spring::Bad(c) => *c,
-                _ => 0,
-            })
-            .sum();
-
-        let expected_ones: usize = total_bad_count - known_bad_count;
-
-        let possibilities = count_choices(
-            &unknown_indexes,
-            &mut Vec::with_capacity(expected_ones),
-            expected_ones,
-            |bad_springs_prefix: &[usize]| self.validate_check_prefix(&bad_springs_prefix),
-            |bad_springs: &[usize]| self.validate_check(&bad_springs),
-        );
-
-        dbg!(possibilities)
+        if has_extra_good_or_unknown_to_trim {
+            self.springs.pop();
+        }
     }
 
-    fn possibilities_count_alt(&self) -> usize {
-        let working_row = WorkingSpringRow {
-            springs: &self.springs,
-            check: &self.check,
+    fn possibilities_count_inner(
+        &self,
+        springs_offset: usize,
+        check_offset: usize,
+        cache: &mut BTreeMap<(usize, usize), usize>,
+    ) -> usize {
+        if let Some(cached) = cache.get(&(springs_offset, check_offset)) {
+            return *cached;
         };
-        dbg!(working_row.possibilities_count())
-    }
-}
-
-impl Spring {
-    fn parser(input: &str) -> IResult<&str, Self> {
-        alt((
-            value(Spring::Good, many1_count(char('.'))),
-            map(many1_count(char('#')), Spring::Bad),
-            value(Spring::Unknown, char('?')),
-        ))(input)
-    }
-}
-
-fn count_choices(
-    possibilities: &[usize],
-    prefix: &mut Vec<usize>,
-    pulls: usize,
-    prefix_filter: impl Fn(&[usize]) -> (bool, usize) + Copy,
-    filter: impl Fn(&[usize]) -> bool + Copy,
-) -> usize {
-    let mut res = 0;
-    if pulls == 0 {
-        res = 1;
-    } else if pulls == 1 {
-        for suffix in possibilities {
-            prefix.push(*suffix);
-            if filter(&prefix) {
-                res += 1;
-            }
-            prefix.pop();
-        }
-    } else {
-        for (i, suffix) in possibilities.iter().enumerate() {
-            prefix.push(*suffix);
-            let (prefix_could_be_valid, prefix_filter_read_count) = prefix_filter(&prefix);
-            if prefix_could_be_valid {
-                res += count_choices(
-                    &possibilities[i + 1..],
-                    prefix,
-                    pulls - 1,
-                    prefix_filter,
-                    filter,
-                );
-            }
-            prefix.pop();
-
-            if !prefix_could_be_valid && prefix_filter_read_count < prefix.len() {
-                break;
-            }
-        }
-    }
-    res
-}
 
-impl<'a> WorkingSpringRow<'a> {
-    fn possibilities_count(&self) -> usize {
-        // this now definitely begins and ends with unknowns
-        let optimized = match self.optimize() {
+        let (springs_offset, check_offset) = match self.optimize_head(springs_offset, check_offset)
+        {
             Some(optimized) => optimized,
             None => {
                 return 0;
             }
         };
 
-        if optimized.check.len() == 1
-            && optimized
-                .springs
-                .iter()
-                .all(|s| !matches!(s, Spring::Bad(_)))
-        {
+        let springs = &self.springs[springs_offset..];
+        let check = &self.check[check_offset..];
+
+        if check.len() == 1 && springs.iter().all(|s| !matches!(s, Spring::Bad(_))) {
             let mut contigous_unknowns = Vec::new();
             let mut next_contiguous_unknown = 0;
             let mut is_in_unknown = false;
-            for spring in optimized.springs.iter() {
+            for spring in springs.iter() {
                 match spring {
                     Spring::Unknown => {
                         next_contiguous_unknown += 1;
@@ -320,8 +196,8 @@ impl<'a> WorkingSpringRow<'a> {
             return contigous_unknowns
                 .iter()
                 .map(|region| {
-                    if *region >= optimized.check[0] as usize {
-                        region - optimized.check[0] as usize + 1
+                    if *region >= check[0] as usize {
+                        region - check[0] as usize + 1
                     } else {
                         0
                     }
@@ -329,12 +205,8 @@ impl<'a> WorkingSpringRow<'a> {
                 .sum();
         }
 
-        if optimized.check.len() == 0 {
-            if optimized
-                .springs
-                .iter()
-                .any(|s| matches!(s, Spring::Bad(_)))
-            {
+        if check.len() == 0 {
+            if springs.iter().any(|s| matches!(s, Spring::Bad(_))) {
                 return 0;
             } else {
                 return 1;
@@ -342,27 +214,72 @@ impl<'a> WorkingSpringRow<'a> {
         }
 
         let valid_prefix_possibilities_count =
-            if let Some(without_prefix) = optimized.trim_bad_prefix() {
-                without_prefix.possibilities_count()
+            if let Some((without_prefix_springs_offset, without_prefix_check_offset)) =
+                self.trim_bad_prefix(springs_offset, check_offset)
+            {
+                self.possibilities_count_inner(
+                    without_prefix_springs_offset,
+                    without_prefix_check_offset,
+                    cache,
+                )
             } else {
                 0
             };
 
-        let non_prefix_possibilities_count = {
-            let subset = WorkingSpringRow {
-                springs: &optimized.springs[1..],
-                check: &optimized.check,
-            };
-            subset.possibilities_count()
-        };
+        let non_prefix_possibilities_count =
+            self.possibilities_count_inner(springs_offset + 1, check_offset, cache);
+
+        let total_possibilities = valid_prefix_possibilities_count + non_prefix_possibilities_count;
+
+        cache.insert((springs_offset, check_offset), total_possibilities);
+
+        total_possibilities
+    }
+
+    fn optimize_head(
+        &self,
+        mut springs_offset: usize,
+        mut check_offset: usize,
+    ) -> Option<(usize, usize)> {
+        while springs_offset < self.springs.len() && check_offset < self.check.len() {
+            match self.springs[springs_offset] {
+                Spring::Good => {
+                    springs_offset += 1;
+                }
+                Spring::Bad(s) if s > self.check[check_offset] as usize => {
+                    return None;
+                }
+                Spring::Bad(s) if s <= self.check[check_offset] as usize => {
+                    match self.trim_bad_prefix(springs_offset, check_offset) {
+                        Some((new_springs, new_check)) => {
+                            springs_offset = new_springs;
+                            check_offset = new_check;
+                        }
+                        None => return None,
+                    };
+                }
+                Spring::Bad(_) => unreachable!(),
+                Spring::Unknown => {
+                    break;
+                }
+            }
+        }
 
-        valid_prefix_possibilities_count + non_prefix_possibilities_count
+        if springs_offset == self.springs.len() && check_offset != self.check.len() {
+            None
+        } else {
+            Some((springs_offset, check_offset))
+        }
     }
 
-    fn trim_bad_prefix(&self) -> Option<WorkingSpringRow<'a>> {
-        let first_check = self.check[0] as usize;
+    fn trim_bad_prefix(
+        &self,
+        springs_offset: usize,
+        check_offset: usize,
+    ) -> Option<(usize, usize)> {
+        let first_check = self.check[check_offset] as usize;
         let mut check_i = 0;
-        let mut springs_i = 0;
+        let mut springs_i = springs_offset;
         while check_i < first_check {
             if springs_i >= self.springs.len() {
                 return None;
@@ -394,106 +311,19 @@ impl<'a> WorkingSpringRow<'a> {
             } else {
                 springs_i
             };
-            Some(WorkingSpringRow {
-                springs: &self.springs[new_springs_i..],
-                check: &self.check[1..],
-            })
-        } else {
-            None
-        }
-    }
-
-    fn trim_bad_suffix(&self) -> Option<WorkingSpringRow<'a>> {
-        let last_check = self.check[self.check.len() - 1] as usize;
-        let mut check_i = 0;
-        let mut new_springs_len = self.springs.len();
-        while check_i < last_check {
-            if new_springs_len == 0 {
-                return None;
-            }
-            match self.springs[new_springs_len - 1] {
-                Spring::Unknown => {
-                    check_i += 1;
-                }
-                Spring::Bad(inc) => {
-                    check_i += inc;
-                }
-                Spring::Good => {
-                    return None;
-                }
-            }
-            new_springs_len -= 1;
-        }
-        let is_boundary = new_springs_len == 0;
-        let has_extra_good_or_unknown_to_trim =
-            !is_boundary && !matches!(self.springs[new_springs_len - 1], Spring::Bad(_));
-        let valid_suffix =
-            check_i == last_check && (new_springs_len == 0 || has_extra_good_or_unknown_to_trim);
-        if valid_suffix {
-            let new_springs_len = if has_extra_good_or_unknown_to_trim {
-                new_springs_len - 1
-            } else {
-                new_springs_len
-            };
-            Some(WorkingSpringRow {
-                springs: &self.springs[..new_springs_len],
-                check: &self.check[..self.check.len() - 1],
-            })
+            Some((new_springs_i, check_offset + 1))
         } else {
             None
         }
     }
+}
 
-    fn optimize(&'a self) -> Option<WorkingSpringRow<'a>> {
-        let mut optimized = self.clone();
-        while optimized.springs.len() > 0 && optimized.check.len() > 0 {
-            match optimized.springs[0] {
-                Spring::Good => {
-                    optimized.springs = &optimized.springs[1..];
-                }
-                Spring::Bad(s) if s > optimized.check[0] as usize => {
-                    return None;
-                }
-                Spring::Bad(s) if s <= optimized.check[0] as usize => {
-                    let trimmed = match optimized.trim_bad_prefix() {
-                        Some(opt) => opt,
-                        None => return None,
-                    };
-                    optimized = trimmed;
-                }
-                Spring::Bad(_) => unreachable!(),
-                Spring::Unknown => {
-                    break;
-                }
-            }
-        }
-
-        while optimized.springs.len() > 0 && optimized.check.len() > 0 {
-            match optimized.springs[optimized.springs.len() - 1] {
-                Spring::Good => {
-                    optimized.springs = &optimized.springs[..optimized.springs.len() - 1];
-                }
-                Spring::Bad(s) if s > optimized.check[optimized.check.len() - 1] as usize => {
-                    return None;
-                }
-                Spring::Bad(s) if s <= optimized.check[optimized.check.len() - 1] as usize => {
-                    let trimmed = match optimized.trim_bad_suffix() {
-                        Some(opt) => opt,
-                        None => return None,
-                    };
-                    optimized = trimmed;
-                }
-                Spring::Bad(_) => unreachable!(),
-                Spring::Unknown => {
-                    break;
-                }
-            }
-        }
-
-        if optimized.springs.len() == 0 && optimized.check.len() != 0 {
-            None
-        } else {
-            Some(optimized)
-        }
+impl Spring {
+    fn parser(input: &str) -> IResult<&str, Self> {
+        alt((
+            value(Spring::Good, many1_count(char('.'))),
+            map(many1_count(char('#')), Spring::Bad),
+            value(Spring::Unknown, char('?')),
+        ))(input)
     }
 }