adventofcode-2020/day-7/src/main.rs

use std::collections::HashMap;
use std::fs;

use lazy_static::lazy_static;
use regex::Regex;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("input")?;
    let rules = parse_rules(&input)?;

    // Part 1
    println!("{}", count_containing_bags(&rules, "shiny gold"));

    // Part 2
    println!("{}", count_contained_bags(&rules, "shiny gold"));

    Ok(())
}

fn parse_rule(input: &str) -> Result<HashMap<String, usize>, std::num::ParseIntError> {
    // Rules look something like:
    //
    //     <value> <key> bags?, <value> <key> bags?.
    //
    // or:
    //
    //     no other bags.

    lazy_static! {
        static ref RULE_RE: Regex = Regex::new(r"(\d+) (.+?) bags?").expect("Should compile");
    };

    if let "no other bags." = input {
        Ok(HashMap::new())
    } else {
        RULE_RE
            .captures_iter(input)
            .map(|captures| Ok((captures[2].to_string(), captures[1].parse::<usize>()?)))
            .collect()
    }
}

fn parse_rules(input: &str) -> Result<HashMap<&str, HashMap<String, usize>>, String> {
    input
        .lines()
        .map(|line| {
            // Each line looks something like:
            //
            //     <key> bags contain <rule>
            //
            // So we simply split on " bags contain " and pass the
            // rest to the rule parsing function.

            // Note this will be a lot easier once str::split_once is stable
            let (key, rule) = {
                let temp: Vec<&str> = line.split(" bags contain ").collect();
                if temp.len() != 2 {
                    Err(format!("Invalid format for bag rule: {}", line))
                } else {
                    Ok((temp[0], temp[1]))
                }
            }?;

            Ok((
                key,
                parse_rule(rule).map_err(|e| format!("Could not parse rule '{}': {}", rule, e))?,
            ))
        })
        .collect()
}

fn count_containing_bags(rules: &HashMap<&str, HashMap<String, usize>>, color: &str) -> usize {
    fn contains_recursively(
        rules: &HashMap<&str, HashMap<String, usize>>,
        rule: &HashMap<String, usize>,
        color: &str,
    ) -> bool {
        // A bag contains the given color recursively, if either it
        // contains the color, or any bags it contains contain the bag
        // recursively
        rule.contains_key(color)
            || rule.keys().any(|bag| {
                if let Some(rule) = rules.get(bag.as_str()) {
                    contains_recursively(rules, rule, color)
                } else {
                    // The bag should always be in our index, but in
                    // case it's not, we treat it as if it contains no other bags.
                    false
                }
            })
    }

    // Count from how many bags the given bag is reachable
    rules
        .iter()
        .filter(|(_, rule)| contains_recursively(rules, rule, color))
        .count()
}

fn count_contained_bags(rules: &HashMap<&str, HashMap<String, usize>>, color: &str) -> usize {
    if let Some(rule) = rules.get(color) {
        // Count how many bags are reachable from this bag; except
        // multiply the numbers by the number of each bag contained
        rule.iter()
            .map(|(contained, number)| number + number * count_contained_bags(rules, contained))
            .sum()
    } else {
        // The bag should always be in our index, but in case it's
        // not, we treat it as if it contains 0 other bags.
        0
    }
}

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

    #[test]
    fn test_simple() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "
            light red bags contain 1 bright white bag, 2 muted yellow bags.
            dark orange bags contain 3 bright white bags, 4 muted yellow bags.
            bright white bags contain 1 shiny gold bag.
            muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
            shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
            dark olive bags contain 3 faded blue bags, 4 dotted black bags.
            vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
            faded blue bags contain no other bags.
            dotted black bags contain no other bags.
            "
        );

        let rules = parse_rules(&input)?;
        assert_eq!(count_containing_bags(&rules, "shiny gold"), 4);

        Ok(())
    }

    #[test]
    fn test_simple2() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "
            light red bags contain 1 bright white bag, 2 muted yellow bags.
            dark orange bags contain 3 bright white bags, 4 muted yellow bags.
            bright white bags contain 1 shiny gold bag.
            muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
            shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
            dark olive bags contain 3 faded blue bags, 4 dotted black bags.
            vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
            faded blue bags contain no other bags.
            dotted black bags contain no other bags.
            "
        );

        let rules = parse_rules(&input)?;
        assert_eq!(count_contained_bags(&rules, "shiny gold"), 32);

        Ok(())
    }

    #[test]
    fn test_simple3() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "
            shiny gold bags contain 2 dark red bags.
            dark red bags contain 2 dark orange bags.
            dark orange bags contain 2 dark yellow bags.
            dark yellow bags contain 2 dark green bags.
            dark green bags contain 2 dark blue bags.
            dark blue bags contain 2 dark violet bags.
            dark violet bags contain no other bags.
            "
        );

        let rules = parse_rules(&input)?;
        assert_eq!(count_contained_bags(&rules, "shiny gold"), 126);

        Ok(())
    }
}