aoc-2020/d07/main.py

83 lines
2.1 KiB
Python
Executable File

#! /usr/bin/env python3
import itertools
import re
from collections import namedtuple
line_matcher = re.compile(r"(\w+ \w+) bags contain ([\w+ ,]+)+.")
child_matcher = re.compile(r"([0-9]+) (\w+ \w+) bag[s]{0,1}")
BagRule = namedtuple("BagRule", ["parent_bag", "bag_counts"])
def parse_line(line):
result = line_matcher.match(line)
if result is None:
# print(f"{line}: No match")
return None
else:
parent_bag = result.group(1)
bag_counts = {}
# print(result.groups())
for bags in result.group(2).split(","):
if bags == "no other bags":
break
child_result = child_matcher.search(bags.strip())
# print(bags.strip(), child_result.groups())
if child_result is None:
print(f"Invalid child bags: {bags}")
bag_counts[child_result.group(2)] = int(child_result.group(1))
return BagRule(parent_bag, bag_counts)
def get_direct_carriers(contained_in, bags):
return set(itertools.chain.from_iterable((
contained_in.get(bag, [])
for bag in bags
)))
def part1():
contained_in = {}
with open("input.txt") as f:
for line in f:
rule = parse_line(line)
for bag in rule.bag_counts.keys():
contained_in.setdefault(bag, []).append(rule.parent_bag)
my_bag = "shiny gold"
holding_bags = set()
last_level = set((my_bag,))
while last_level:
last_level = get_direct_carriers(contained_in, last_level)
holding_bags |= last_level
print(f"Holding bags: {len(holding_bags)}")
def count_children(contents, bag):
rule = contents.get(bag)
if rule is None:
return 0
return sum(
count * (1 + count_children(contents, child_bag))
for child_bag, count in rule.bag_counts.items()
)
def part2():
contents = {}
with open("input.txt") as f:
for line in f:
rule = parse_line(line)
contents[rule.parent_bag] = rule
children = count_children(contents, "shiny gold")
print(f"Holds {children} bags")
if __name__ == "__main__":
part1()
part2()