Sean McLemon | Advent of Code

2020-12-07 - Handy Haversacks

(original .ipynb)

Day 7 puzzle input is a list of descriptions of coloured bags and what colours of bags they can contain (mine is here). Part 1 involve s finding all the bag colours that may directly or indirectly contain a "shiny gold" bag. Part 2 involves finding the total number of bags that are inside a "shiny gold" bag.

import re

def parse_bag_str(bag_str):
    m = re.search("(\d+ |no )?(\w+) (\w+) bags?.?", bag_str)
    return (m.group(1), m.group(2), m.group(3))
    
class Bag:
    def __init__(self, bag_str, parent=None):
        quantity, adjective1, adjective2 = parse_bag_str(bag_str)
        self.quantity = int(quantity) if quantity else 0
        self.color = adjective1 + " " + adjective2        
        self.parent = parent
        self.children = []
        
    def add_child(self, child):
        if child.color != "no other":
            self.children.append(child)

def parse_line(line):
    parent_str, children_str = line.split("contain")
    parent = Bag(parent_str)
    
    for child_str in children_str.split(", "):
        child = Bag(child_str, parent)
        parent.add_child(child)
        
    return parent

def contains_color(bag, bag_lookup, color="shiny gold"):
    if not bag.children:
        return False
    
    directly_contains_color = color in [c.color for c in bag.children]
    indirectly_contains_color = any(
        contains_color(bag_lookup[c.color], bag_lookup, color) 
        for c in bag.children)
    
    return directly_contains_color or indirectly_contains_color

def parse_input_str(input_str):
    bag_lookup = {}    
    all_bags = []
    for line in input_str.split("\n"):
        bag = parse_line(line)
        all_bags += [bag] + bag.children
        bag_lookup[bag.color] = bag
    return (all_bags, bag_lookup)

test_input_str = """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."""

puzzle_input_str = open("puzzle_input/day7.txt", "r").read().strip()

def part1(input_str):
    all_bags, bag_lookup = parse_input_str(input_str)
    
    holding_target_color = set()
    for bag in all_bags:
        if contains_color(bag, bag_lookup):
            holding_target_color.add(bag.color)        
    
    return len(holding_target_color)

assert 4 == part1(test_input_str)
print(part1(puzzle_input_str))

378

test_input_str_2 = """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."""

def count_bag_totals(bag, bag_lookup):
    if len(bag.children) == 0:
        return 0
    
    return sum([c.quantity + 
                (c.quantity * count_bag_totals(bag_lookup[c.color], bag_lookup)) 
                for c in bag.children])

def part2(input_str):
    all_bags, bag_lookup = parse_input_str(input_str)
    return count_bag_totals(bag_lookup["shiny gold"], bag_lookup)

assert 126 == part2(test_input_str_2)
print(part2(puzzle_input_str))

27526