aoc2020/day07.py

82 lines
2.3 KiB
Python

import aoclib
import re
DAY = 7
TEST_SOLUTION_PART1 = 4
TEST_SOLUTION_PART2 = 32
MY_BAG = "shiny gold"
content_re = re.compile(r"(\d+) (\w+ \w+) bag.*")
def getBagContent(bag_line):
"""
think of this like a shopping list
"dark gold" needs:
- "posh crimson": 3
- "mirrored lavender": 3
"""
color, contents = bag_line.split(" bags contain ")
return_dict = {color: {}}
if contents == 'no other bags.':
return return_dict
for content in contents.split(", "):
match_groups = re.match(content_re, content)
if not match_groups:
print("ERROR: content '%s' cannot be matched" % content)
return_dict[color][match_groups.group(2)] = int(match_groups.group(1))
return return_dict
def getPossibleBagsForMyColor(color_dict, my_color):
if my_color not in color_dict:
return set()
color_set = set(color_dict[my_color])
for color in color_dict[my_color]:
extend_set = getPossibleBagsForMyColor(color_dict, color)
for extra_color in extend_set:
color_set.add(extra_color)
return color_set
def getBagContentCount(color_dict, my_color):
if my_color not in color_dict or not color_dict[my_color]:
return 0
count = 0
for color in color_dict[my_color]:
sub_count = getBagContentCount(color_dict, color)
count = count + color_dict[my_color][color] + color_dict[my_color][color] * sub_count
return count
def part1(test_mode=False):
my_input = aoclib.getInputAsArray(day=DAY, test=test_mode)
possible_bags = {}
for bag_line in my_input:
content_dict = getBagContent(bag_line)
for bag_color in content_dict:
for content_color in content_dict[bag_color]:
if content_color not in possible_bags:
possible_bags[content_color] = [bag_color]
else:
possible_bags[content_color].append(bag_color)
return len(getPossibleBagsForMyColor(possible_bags, MY_BAG))
def part2(test_mode=False):
my_input = aoclib.getInputAsArray(day=DAY, test=test_mode)
big_content_dict = {}
for bag_line in my_input:
content_dict = getBagContent(bag_line)
for color in content_dict:
big_content_dict[color] = content_dict[color]
return getBagContentCount(big_content_dict, MY_BAG)