57 lines
1.1 KiB
Python
57 lines
1.1 KiB
Python
import re
|
|
|
|
def process_rules(data):
|
|
rules = {}
|
|
for line in data:
|
|
color = re.search("(\w+ \w+) bags contain", line).group(1)
|
|
rule = re.findall("(\d+) (\w+ \w+)", line)
|
|
|
|
rules[color] = []
|
|
for bag in rule:
|
|
rules[color].append((bag[1], int(bag[0])))
|
|
|
|
return rules
|
|
|
|
|
|
def find_in_bags(color, rules):
|
|
count = 0
|
|
|
|
for bag in rules:
|
|
discovered = []
|
|
discovered.extend(rules[bag])
|
|
|
|
while discovered:
|
|
x, _ = discovered.pop(0)
|
|
if x == 'shiny gold':
|
|
count += 1
|
|
break
|
|
discovered.extend(rules[x])
|
|
|
|
return count
|
|
|
|
def count_bags(bag, rules):
|
|
discovered = [bag]
|
|
total = 0
|
|
|
|
while discovered:
|
|
bag, amount = discovered.pop(0)
|
|
for x, x_amt in rules[bag]:
|
|
discovered.append((x, amount * x_amt))
|
|
total += amount
|
|
|
|
return total - 1
|
|
|
|
|
|
def solve_a(rules):
|
|
return find_in_bags("shiny gold", rules)
|
|
|
|
def solve_b(rules):
|
|
return count_bags(("shiny gold", 1), rules)
|
|
|
|
|
|
data = [line.strip() for line in open('input')]
|
|
rules = process_rules(data)
|
|
print(solve_a(rules))
|
|
print(solve_b(rules))
|
|
|