aoc2020/7.py

54 lines
1.3 KiB
Python

import queue
with open("7.txt") as f:
edges = dict()
edges2 = dict()
import re
pat = re.compile("""
^(.*)\ bags?\ contain\
(?:
(?:no\ other\ bags)
| (?: (\d+)\ (.*)\ bags? (,\ (\d+)\ (.*)\ bags?)* )
)
.$
""", flags=re.VERBOSE)
for line in f:
a, b = line.split("contain")
outer = a.strip().rstrip("bags").strip()
if outer not in edges2:
edges2[outer] = dict()
inner = b.strip().rstrip(".").split(", ")
# print(inner)
if inner == ["no other bags"]: continue
for n in inner:
num = int(re.match("^(\d+)\s.*", n).groups()[0])
color = re.sub("^(\d+)\s", "", n)
color = re.sub("\sbags?$", "", color)
# print(color)
edges2[outer][color] = num
if color not in edges:
edges[color] = set()
edges[color].add(outer)
# print(edges)
found = set()
q = queue.Queue()
q.put("shiny gold")
while not q.empty():
c = q.get()
r = edges.get(c, set())
# print(c, r)
for c1 in r:
if c1 not in found:
q.put(c1)
found.add(c1)
print(len(found))
print(edges2)
def visit(color):
r = edges2.get(color, set())
c = 1
for col, n in r.items():
c += n * visit(col)
# print(f"{color} contains {c} other bags")
return c
print(visit("shiny gold") - 1)