Day 12, attempt at part 2

2 files changed, 388 insertions, 60 deletions

diff --git a/2023/Cargo.toml b/2023/Cargo.toml
index d7b5bda..87e205f 100644
--- a/2023/Cargo.toml
+++ b/2023/Cargo.toml
@@ -10,4 +10,4 @@ nom = "7.1.3"
 thiserror = "1.0.50"
 
 [profile.release]
-# debug = true
+debug = true
diff --git a/2023/src/bin/day_12.rs b/2023/src/bin/day_12.rs
index a2e83be..7485aae 100644
--- a/2023/src/bin/day_12.rs
+++ b/2023/src/bin/day_12.rs
@@ -2,7 +2,7 @@ use nom::{
     branch::alt,
     character::complete::{char, line_ending, space1, u32},
     combinator::{map, value},
-    multi::{many1, separated_list1},
+    multi::{many1, many1_count, separated_list1},
     sequence::separated_pair,
     IResult,
 };
@@ -24,14 +24,21 @@ struct SpringField(Vec<SpringRow>);
 
 #[derive(Debug)]
 struct SpringRow {
-    springs: Vec<Option<Spring>>,
+    springs: Vec<Spring>,
     check: Vec<u32>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum Spring {
     Good,
-    Bad,
+    Bad(usize),
+    Unknown,
+}
+
+#[derive(Debug, Clone)]
+struct WorkingSpringRow<'a> {
+    springs: &'a [Spring],
+    check: &'a [u32],
 }
 
 impl SpringField {
@@ -40,7 +47,7 @@ impl SpringField {
     }
 
     fn possibilities_sum(&self) -> usize {
-        self.0.iter().map(|r| r.possibilities_count()).sum()
+        self.0.iter().map(|r| r.possibilities_count_alt()).sum()
     }
 
     fn expand(&self) -> SpringField {
@@ -68,57 +75,116 @@ impl SpringRow {
 
         for _ in 0..5 {
             expanded.springs.append(&mut self.springs.clone());
-            expanded.springs.push(None);
+            expanded.springs.push(Spring::Unknown);
             expanded.check.append(&mut self.check.clone());
         }
 
-        // should only have the extra None between, not at the very end.
+        // should only have the extra Unknown between, not at the very end.
         expanded.springs.pop();
 
         expanded
     }
 
-    fn validate_check(&self, springs: &[Spring]) -> bool {
-        let mut current_check_index = 0;
+    /// 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 spring in springs {
-            match spring {
-                Spring::Good => {
-                    if in_bad_lands {
-                        let valid = self
-                            .check
-                            .get(current_check_index)
-                            .map_or(false, |expected| *expected == current_count);
-                        if !valid {
-                            return false;
-                        }
-                        current_count = 0;
+        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);
+                    }
+
+                    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)
+                }
+            };
 
-                        current_check_index += 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, bad_spring_read_count);
                     }
+
+                    current_count = 0;
+                    current_check = check_iter.next();
                     in_bad_lands = false;
                 }
-                Spring::Bad => {
-                    current_count += 1;
-                    in_bad_lands = true;
+            }
+        }
+
+        (true, bad_spring_read_count)
+    }
+
+    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),
+                Spring::Unknown => {
+                    let bad_check = current_bad_spring == Some(&i);
+                    if bad_check {
+                        current_bad_spring = bad_spring_iter.next();
+                    }
+                    (bad_check, 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;
                 }
             }
         }
 
         if in_bad_lands {
-            let valid = self
-                .check
-                .get(current_check_index)
-                .map_or(false, |expected| *expected == current_count);
+            let valid = current_check.map_or(false, |expected| *expected == current_count as u32);
             if !valid {
                 return false;
             }
-            current_check_index += 1;
+            current_check = check_iter.next();
         }
 
-        self.check.len() == current_check_index
+        current_check == None
     }
 
     fn possibilities_count(&self) -> usize {
@@ -126,46 +192,308 @@ impl SpringRow {
             .springs
             .iter()
             .enumerate()
-            .filter_map(|(i, s)| s.is_none().then_some(i))
+            .filter_map(|(i, s)| (s == &Spring::Unknown).then_some(i))
             .collect();
 
-        assert_ne!(unknown_indexes.len(), 0);
-        let max_possibilities = 2_usize.pow(unknown_indexes.len().try_into().unwrap());
-
-        let mut spring_guess: Vec<Spring> = self
+        let total_bad_count: usize = self.check.iter().sum::<u32>() as usize;
+        let known_bad_count: usize = self
             .springs
             .iter()
-            .map(|s| s.unwrap_or(Spring::Bad))
-            .collect();
-        let unknown_indexes: Vec<usize> = self
-            .springs
-            .iter()
-            .enumerate()
-            .filter_map(|(i, s)| s.is_none().then_some(i))
-            .collect();
-        let possibilities = (0..max_possibilities)
-            .filter(|i| {
-                for (bit_index, spring_index) in unknown_indexes.iter().enumerate() {
-                    spring_guess[*spring_index] = if i & (1 << bit_index) != 0 {
-                        Spring::Bad
-                    } else {
-                        Spring::Good
-                    };
-                }
-                self.validate_check(&spring_guess)
+            .map(|s| match s {
+                Spring::Bad(c) => *c,
+                _ => 0,
             })
-            .count();
+            .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),
+        );
 
-        possibilities
+        dbg!(possibilities)
+    }
+
+    fn possibilities_count_alt(&self) -> usize {
+        let working_row = WorkingSpringRow {
+            springs: &self.springs,
+            check: &self.check,
+        };
+        dbg!(working_row.possibilities_count())
     }
 }
 
 impl Spring {
-    fn parser(input: &str) -> IResult<&str, Option<Self>> {
+    fn parser(input: &str) -> IResult<&str, Self> {
         alt((
-            value(Some(Spring::Good), char('.')),
-            value(Some(Spring::Bad), char('#')),
-            value(None, char('?')),
+            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() {
+            Some(optimized) => optimized,
+            None => {
+                return 0;
+            }
+        };
+
+        if optimized.check.len() == 1
+            && optimized
+                .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() {
+                match spring {
+                    Spring::Unknown => {
+                        next_contiguous_unknown += 1;
+                        is_in_unknown = true;
+                    }
+                    Spring::Good => {
+                        if is_in_unknown {
+                            contigous_unknowns.push(next_contiguous_unknown);
+                            next_contiguous_unknown = 0;
+                            is_in_unknown = false;
+                        }
+                    }
+                    Spring::Bad(_) => unreachable!(),
+                }
+            }
+            if is_in_unknown {
+                contigous_unknowns.push(next_contiguous_unknown);
+            }
+
+            return contigous_unknowns
+                .iter()
+                .map(|region| {
+                    if *region >= optimized.check[0] as usize {
+                        region - optimized.check[0] as usize + 1
+                    } else {
+                        0
+                    }
+                })
+                .sum();
+        }
+
+        if optimized.check.len() == 0 {
+            if optimized
+                .springs
+                .iter()
+                .any(|s| matches!(s, Spring::Bad(_)))
+            {
+                return 0;
+            } else {
+                return 1;
+            }
+        }
+
+        let valid_prefix_possibilities_count =
+            if let Some(without_prefix) = optimized.trim_bad_prefix() {
+                without_prefix.possibilities_count()
+            } else {
+                0
+            };
+
+        let non_prefix_possibilities_count = {
+            let subset = WorkingSpringRow {
+                springs: &optimized.springs[1..],
+                check: &optimized.check,
+            };
+            subset.possibilities_count()
+        };
+
+        valid_prefix_possibilities_count + non_prefix_possibilities_count
+    }
+
+    fn trim_bad_prefix(&self) -> Option<WorkingSpringRow<'a>> {
+        let first_check = self.check[0] as usize;
+        let mut check_i = 0;
+        let mut springs_i = 0;
+        while check_i < first_check {
+            if springs_i >= self.springs.len() {
+                return None;
+            }
+
+            match self.springs[springs_i] {
+                Spring::Unknown => {
+                    check_i += 1;
+                }
+                Spring::Bad(inc) => {
+                    check_i += inc;
+                }
+                Spring::Good => {
+                    return None;
+                }
+            }
+            springs_i += 1;
+        }
+
+        let is_boundary = springs_i == self.springs.len();
+        let has_extra_good_or_unknown_to_trim =
+            !is_boundary && !matches!(self.springs[springs_i], Spring::Bad(_));
+        let valid_prefix =
+            check_i == first_check && (is_boundary || has_extra_good_or_unknown_to_trim);
+
+        if valid_prefix {
+            let new_springs_i = if has_extra_good_or_unknown_to_trim {
+                springs_i + 1
+            } 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],
+            })
+        } 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)
+        }
+    }
+}