07.koto

# https://adventofcode.com/2020/day/7

parse_bag_rules = |input|
  input
    .lines()
    .each |line|
      # e.g. light red bags contain 1 bright white bag, 2 muted yellow bags.
      # e.g.  faded blue bags contain no other bags.
      words = line.split " "
      bag = "${words.next()} ${words.next()}"
      words.next() # bags
      words.next() # contain
      contained_bags = {}
      loop
        match words.next()
          "no" or null then 
            break bag, contained_bags
          quantity then
            contained_bag = "${words.next()} ${words.next()}"
            contained_bags.insert contained_bag, quantity.to_number()
            words.next() # bag(s)
    .to_map()

count_bag_options = |target_bag, rules|
  bag_contains_target_bag = |bag|
    rules
      .get(bag)
      .any |(option, _)| option == target_bag or bag_contains_target_bag option

  rules
    .keep |(bag, options)| bag != target_bag and bag_contains_target_bag bag
    .count()

count_contained_bags = |bag, rules|
  rules
    .get bag
    .each |(contained_bag, count)|
      count + count * count_contained_bags contained_bag, rules
    .sum()

@main = ||
  rules = io.extend_path koto.script_dir, "input", "07"
    >> io.read_to_string
    >> parse_bag_rules

  print "Part one: ${count_bag_options "shiny gold", rules}" # 272
  print "Part two: ${count_contained_bags "shiny gold", rules}" # 172246

@tests =
  @pre_test: ||
    self.example_1 = "\
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: ||
    rules = parse_bag_rules self.example_1
    assert_eq (count_bag_options "shiny gold", rules), 4

  @test part_two_example_1: ||
    rules = parse_bag_rules self.example_1
    assert_eq (count_contained_bags "shiny gold", rules), 32

  @test part_two_example_2: ||
    rules = parse_bag_rules "\
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.
"
    assert_eq (count_contained_bags "shiny gold", rules), 126