Complete day 7

day-7/src/main.rs

@@ -0,0 +1,182 @@
+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(())
+    }
+}