aoc 2020 day 7 python solution

day7.py

import re


def parse(filename):
    bags = {}
    inverted_bags = {}

    with open(filename) as in_f:
        for line in in_f.readlines():
            color, _, rest = map(lambda s: s.strip(), line.partition("bags contain "))
            bags[color] = {}
            if not rest.startswith("no other"):
                parts = re.sub(r'bag[s\.]*', ' ', rest).split(",")
                parts = list(map(lambda s: s.strip().partition(" "), parts))
                for part in parts:
                    ccolor, ccount = part[2], int(part[0])
                    bags[color][ccolor] = ccount
                    if ccolor in inverted_bags:
                        inverted_bags[ccolor].add(color)
                    else:
                        inverted_bags[ccolor] = {color}
    return (bags, inverted_bags)


def part1(inverted_bags, colors, found = set()):
    if len(colors) != 0:
        for color in colors:
            found = found.union(inverted_bags[color])
            inverted_bags.pop(color)
        colors = set(inverted_bags.keys()).intersection(found)
        return part1(inverted_bags, colors, found)
    return len(found)


def part2(bags, color, current=1):
    el = bags[color]
    if len(el) != 0:
        return current + current * sum(part2(bags, key, el[key]) for key in el)
    else:
        return current


(bags, inverted_bags) = parse("day7.input")
print("part1 = {}".format(part1(inverted_bags, {"shiny gold"})))
print("part2 = {}".format(part2(bags, "shiny gold") - 1))