AdventOfCode2020/07/main.py

import tqdm

class Bag:

    def __init__(self,color):
        self.color = color
        self.capacity = []

    def __str__(self):
        return self.color + str(self.capacity)
    
    def addCapacity(self, amount, color):
        self.capacity.append([amount,color])
    
    def getCapacity(self):
        return self.capacity

    def getColor(self):
        return self.color

    def isColor(self, color):
        if color == self.color:
            return True
        else:
            return False

    def canHold(self, color):
        for x in self.capacity:
            print('['+self.color+'] check '+color+' in '+str(x))
            if color in x:
                return True
            else:
                return False


file = open('input.txt','r')
#file = open('test.txt','r')

rules = file.readlines()
bags = []

def getBag(color):
    found = False
    for b in bags:
        if b.isColor(color):
            found = True
            return b
    if not found:
        b = Bag(color)
        bags.append(b)
        return b

for r in rules:
    r = r.replace('.','')
    r = r.replace('\n','')
    par = r.split(' bags contain')
    b = getBag(par[0])
    if par[1] == ' no other bags':
        continue
    cap = par[1].split(',')
    for c in cap:
        c = c.strip()
        amount = c.split(' ')[0]
        color = c.replace(str(amount)+' ','').replace(' bags','').replace(' bag','')
        b.addCapacity(amount,color)
    
sum = 0

def getToColor(target,start):
    if target == start:
        return True
    
    bag = getBag(start)
    for b in bag.getCapacity():
        if getToColor(target,getBag(b[1]).getColor()):
            return True
    return False

def countSubBags(color):
    sum = 0
    bag = getBag(color)
    if len(bag.getCapacity()) == 0:
        return 1
    for b in bag.getCapacity():
        #print('['+color+'] '+ b[0] + ' * ' + str(countSubBags(b[1])))
        sum += int(b[0]) * countSubBags(b[1])
    print(color + ' ' +str(sum))
    return sum + 1

print(countSubBags('shiny gold') - 1)

exit()
for b in bags:
    if b.getColor() == 'shiny gold':
        continue
    if getToColor('shiny gold',b.getColor()):
            sum += 1

print(str(sum))