import sys | |
sys.path.append("..") | |
from usefulDecorators import * | |
#Creates a dictionary where keys are colors of bag (string) and the values are a set of the | |
#colors of bag which directly contain the key color | |
def createBagList(): | |
bagList = dict() | |
with open('input.txt', 'r') as infile: | |
for line in infile: | |
#Split input line container colored and contained colors | |
container, contains = line.split(' bags contain ') | |
#This line is a hack that assumes all numbers are single digit, which they appear to be?? | |
contains = [c.strip('\n. ').replace(' bags', '').replace(' bag', '')[2:] for c in contains.split(',') if 'no other bags' not in c] | |
#print(f"{container}\t\t{contains}") | |
for color in contains: | |
if color not in bagList: | |
bagList[color] = {container} | |
else: | |
bagList[color].add(container) | |
return bagList | |
def findParentColors(color, bagList): | |
if color not in bagList: | |
return set() | |
else: | |
retVal = bagList[color] | |
return retVal.union(*[findParentColors(c, bagList) for c in bagList[color]]) | |
#Creates a dictionary of values where the keys are the container and the values are a list of colors that container can directly contain | |
#Note that this is essentially the opposite of what createBagList() does | |
def createContainedList(): | |
baseList = dict() | |
with open('input.txt', 'r') as infile: | |
for line in infile: | |
container, contains = line.split(' bags contain ') | |
#This line is a hack that assumes all numbers are single digit, which they appear to be?? | |
contains = [c.strip('\n. ').replace(' bags', '').replace(' bag', '') for c in contains.split(',') if 'no other bags' not in c] | |
contains = [(int(c[0]), c[2:]) for c in contains] | |
baseList[container] = contains | |
return baseList | |
def findTotalBags(color, bagList): | |
if color not in bagList: | |
raise ValueError(f"{color} is not in the list of keys of baglist") | |
else: | |
if len(bagList[color]) == 0: return 1 | |
else: return sum([c[0] * findTotalBags(c[1], bagList) for c in bagList[color]]) + 1 | |
if __name__ == '__main__': | |
#Solution to part 1 | |
print(f"Solution to part 1: {len(findParentColors('shiny gold', createBagList()))}") | |
#Solution to part 2 | |
print(f"Solution to part 2: {findTotalBags('shiny gold', createContainedList())-1}") | |
#Have to subtract 1 from the final answer, because findTotalBags includes itself in the results, | |
#which we want for every bag except the 'shiny gold' bag itself | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment