1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use crate::input::Input;
use std::collections::{HashMap, HashSet};

const MAX_DEPTH: u32 = 100;

#[derive(Debug)]
struct BagEntry<'a> {
    amount: u32,
    bag_type: &'a str,
}

fn insert_ancestors<'a>(
    child_to_parent: &'a HashMap<&'a str, Vec<&str>>,
    child_bag_type: &'a str,
    ancestors: &mut HashSet<&'a str>,
    height: u32,
) -> Result<(), String> {
    if height > MAX_DEPTH {
        return Err(format!(
            "Too deep tree (possibly recursive) - bailing at depth {}",
            MAX_DEPTH
        ));
    }

    if let Some(parents) = child_to_parent.get(child_bag_type) {
        ancestors.extend(parents);
        for parent in parents {
            insert_ancestors(child_to_parent, parent, ancestors, height + 1)?;
        }
    }
    Ok(())
}

fn count_total_bags<'a>(
    reactions: &'a HashMap<&'a str, Vec<BagEntry>>,
    bag_type: &'a str,
    depth: u32,
) -> Result<u32, String> {
    if depth > MAX_DEPTH {
        return Err(format!(
            "Too deep tree (possibly recursive) - bailing at depth {}",
            MAX_DEPTH
        ));
    }

    reactions
        .get(bag_type)
        .map_or(Ok(0_u32), |resulting_entries| {
            resulting_entries
                .iter()
                .map(|entry| {
                    count_total_bags(reactions, entry.bag_type, depth + 1)
                        .map(|value| entry.amount * (value + 1))
                })
                .sum::<Result<u32, String>>()
        })
}

pub fn solve(input: &mut Input) -> Result<u32, String> {
    let mut reactions: HashMap<&str, Vec<BagEntry>> = HashMap::new();
    let mut child_to_parent: HashMap<&str, Vec<&str>> = HashMap::new();

    for (line_idx, line) in input
        .text
        .lines()
        .enumerate()
        .filter(|(_line_idx, line)| !line.contains("no other bags"))
    {
        let on_error = || format!("Line {}: Invalid format", line_idx + 1);

        let mut parts = line.split(" bags contain ");
        let from_bag = parts.next().ok_or_else(on_error)?;

        let mut children_entries = Vec::new();
        let to_parts = parts
            .next()
            .ok_or_else(on_error)?
            .strip_suffix('.')
            .ok_or_else(on_error)?;

        for to_part in to_parts.split(", ") {
            let mut amount_and_bag_type = to_part.splitn(2, ' ');
            let amount = amount_and_bag_type
                .next()
                .ok_or_else(on_error)?
                .parse::<u32>()
                .map_err(|_| on_error())?;
            let bag_type: &str = amount_and_bag_type.next().ok_or_else(on_error)?;
            let bag_type = bag_type.rsplitn(2, ' ').nth(1).ok_or_else(on_error)?;

            if input.is_part_one() {
                child_to_parent
                    .entry(bag_type)
                    .or_insert_with(Vec::new)
                    .push(from_bag);
            } else {
                children_entries.push(BagEntry { amount, bag_type });
            }
        }
        reactions.insert(from_bag, children_entries);
    }

    Ok(if input.is_part_one() {
        let mut gold_ancestors: HashSet<&str> = HashSet::new();
        insert_ancestors(&child_to_parent, "shiny gold", &mut gold_ancestors, 0)?;
        gold_ancestors.len() as u32
    } else {
        count_total_bags(&reactions, "shiny gold", 0)?
    })
}

#[test]
pub fn tests() {
    use crate::{test_part_one, test_part_two};

    let example = "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.";
    test_part_one!(example => 4);
    test_part_two!(example => 32);

    let real_input = include_str!("day07_input.txt");
    test_part_one!(real_input => 229);
    test_part_two!(real_input => 6683);
}