79 lines
2.2 KiB
Python
79 lines
2.2 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(", "):
|
|
content_count, content_color = re.match(content_re, content).groups()
|
|
return_dict[color][content_color] = int(content_count)
|
|
|
|
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)
|