# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# Read file
import itertools
import re

is_num = re.compile(r"^\d+$")

def read_file(filename):
    with open(filename) as infile:
        lines = [line.strip().split() for line in infile.readlines()]
        bags = {}
        for line in lines:
            bag = ' '.join(line[:2])
            i = 4
            contents = {}
            while i < len(line) and is_num.match(line[i]):
                contents[' '.join(line[i+1:i+3])] = int(line[i])
                i += 4
            bags[bag] = contents
    return bags


# -

# Part 1
def contains_shiny_gold(bags, bag, memo=None):
    if memo == None:
        memo = {}
    if bag in memo:
        return memo[bag]
    if 'shiny gold' in bags[bag]:
        memo[bag] = True
        return True
    for inner in bags[bag]:
        if contains_shiny_gold(bags, inner, memo):
            memo[bag] = True
            return True
    memo[bag] = False
    return False


# Test part 1
bags = read_file("test01.txt")
sum([contains_shiny_gold(bags, bag) for bag in bags]) == 4

# Solve part 1
bags = read_file("input.txt")
sum([contains_shiny_gold(bags, bag) for bag in bags])


# Part 2
def count_inner_bags(bags, bag, memo=None):
    if memo == None:
        memo = {}
    if bag in memo:
        return memo[bag]
    count = 0
    for inner in bags[bag]:
        count += count_inner_bags(bags, inner, memo) * bags[bag][inner] + bags[bag][inner]
    memo[bag] = count
    return count


# Test part 2
count_inner_bags(read_file("test01.txt"), 'shiny gold') == 32 and \
count_inner_bags(read_file("test02.txt"), 'shiny gold') == 126

# Solve part 2
count_inner_bags(read_file("input.txt"), 'shiny gold')

StarcoderdataJupyterScriptsDedupFiltered