advent-of-code/2023/src/bin/12.rs

advent_of_code::solution!(12);

use std::collections::HashMap;
use std::str::FromStr;

type Cache = HashMap<(Vec<Condition>, Vec<usize>), usize>;

#[derive(Copy, Clone, Eq, PartialEq, Hash)]
enum Condition {
    Unknown,
    Functional,
    Damaged,
}

impl FromStr for Condition {
    type Err = String;
    fn from_str(input: &str) -> Result<Condition, Self::Err> {
        match input {
            "?" => Ok(Condition::Unknown),
            "." => Ok(Condition::Functional),
            "#" => Ok(Condition::Damaged),
            n @ _ => Err(format!("Unknown character '{}'", n)),
        }
    }
}

#[derive(Clone)]
struct Spring {
    list: Vec<Condition>,
    groups: Vec<usize>,
}

fn parse(input: &str) -> Vec<Spring> {
    input
        .lines()
        .filter_map(|line| {
            if !line.is_empty() {
                let parts = line.split(" ").collect::<Vec<&str>>();
                let list = parts[0]
                    .split("")
                    .filter_map(|e| {
                        if e.is_empty() {
                            None
                        } else {
                            Condition::from_str(e).ok()
                        }
                    })
                    .collect();
                let groups = parts[1]
                    .split(",")
                    .map(|e| e.parse::<usize>().unwrap())
                    .collect();
                Some(Spring { list, groups })
            } else {
                None
            }
        })
        .collect()
}

fn calc_solutions(list: &Vec<Condition>, groups: &Vec<usize>, cache: &mut Cache) -> usize {
    if list.is_empty() {
        if groups.is_empty() {
            return 1;
        } else {
            return 0;
        }
    }

    match list[0] {
        Condition::Functional => calc_solutions(&list[1..].to_vec(), groups, cache),
        Condition::Damaged => calc_damaged_solutions(list, groups, cache),
        Condition::Unknown => {
            calc_solutions(&list[1..].to_vec(), groups, cache)
                + calc_damaged_solutions(list, groups, cache)
        }
    }
}

fn calc_damaged_solutions(list: &Vec<Condition>, groups: &Vec<usize>, cache: &mut Cache) -> usize {
    if let Some(&result) = cache.get(&(list.clone(), groups.clone())) {
        return result;
    }

    if groups.is_empty() {
        return 0;
    }

    if list.len() < groups[0] {
        return 0;
    }

    for i in 0..groups[0] {
        if list[i] == Condition::Functional {
            return 0;
        }
    }

    if list.len() == groups[0] {
        if groups.len() == 1 {
            return 1;
        }
        return 0;
    }
    if list[groups[0]] == Condition::Damaged {
        return 0;
    }
    let result = calc_solutions(
        &list[(groups[0] + 1)..].to_vec(),
        &groups[1..].to_vec(),
        cache,
    );
    cache.insert((list.clone(), groups.clone()), result);
    result
}

pub fn part_one(input: &str) -> Option<u32> {
    let mut result: usize = 0;
    for src in parse(input).iter() {
        let mut cache: Cache = HashMap::new();
        result += calc_solutions(&src.list, &src.groups, &mut cache);
    }
    Some(result as u32)
}

pub fn part_two(input: &str) -> Option<u64> {
    part2(&parse(input))
}

fn part2(input: &[Spring]) -> Option<u64> {
    let mut result: usize = 0;
    for src in input.iter() {
        let mut l: Vec<Condition> = Vec::with_capacity(src.groups.len() * 5 + 5);
        let mut g: Vec<usize> = Vec::with_capacity(src.list.len() * 5);
        for index in 0..5 {
            l.append(&mut src.list.clone());
            if index < 4 {
                l.push(Condition::Unknown);
            }
            g.append(&mut src.groups.clone());
        }
        let mut cache: Cache = HashMap::new();
        let s = calc_solutions(&l, &g, &mut cache);
        result += s;
    }
    Some(result.try_into().unwrap())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_part_one() {
        let result = part_one(&advent_of_code::template::read_file("examples", DAY));
        assert_eq!(result, Some(21));
    }

    #[test]
    fn test_part_two() {
        let result = part_two(&advent_of_code::template::read_file("examples", DAY));
        assert_eq!(result, Some(525152));
    }

    #[test]
    fn sample_data_3() {
        let line1 = Spring {
            list: vec![
                Condition::Unknown,
                Condition::Unknown,
                Condition::Unknown,
                Condition::Functional,
                Condition::Damaged,
                Condition::Damaged,
                Condition::Damaged,
            ],
            groups: vec![1, 1, 3],
        };
        assert_eq!(part2(&[line1]), Some(1));

        let line2 = Spring {
            list: vec![
                Condition::Functional,
                Condition::Unknown,
                Condition::Unknown,
                Condition::Functional,
                Condition::Functional,
                Condition::Unknown,
                Condition::Unknown,
                Condition::Functional,
                Condition::Functional,
                Condition::Functional,
                Condition::Unknown,
                Condition::Damaged,
                Condition::Damaged,
                Condition::Functional,
            ],
            groups: vec![1, 1, 3],
        };
        assert_eq!(part2(&[line2]), Some(16384));
    }
}