Using breadth first search!

2020/day_7/handy_haversacks.py

@@ -1,64 +1,52 @@
-class Bag():
-    def __init__(self, color):
-        self.color = color
-        self.children = []
-
-    def append(self, child):
-        self.children.append(child)
-
-    def count(self):
-        if not self.children:
-            return 1
-
-        return sum([child.count() for child in self.children]) + 1
-
-    def __repr__(self):
-        return f"{self.color} -> {self.children}"
-
-
 def process_rules(data):
     rules = {}
-    for rule in data:
-        color, rule = [r.strip() for r in rule.split("contain")]
+    for line in data:
+        line = line.replace(' bags', '').replace(' bag', '').replace('.', '')
+        color, rule = [r.strip() for r in line.split('contain ')]
+        rule = rule.split(', ')
 
-        color = color.replace('bags', '').strip()
-        rule = rule.replace('.', '').split(', ')
-
-        rules[color] = {}
+        rules[color] = []
         for bag in rule:
-            if bag == 'no other bags': break
-
-            amount = int(bag[0])
-            bag = bag[1:].replace('bags', '').replace('bag', '').strip()
-
-            rules[color][bag] = amount
+            if bag == 'no other': break
+            rules[color].append((bag[2:], int(bag[0])))
 
     return rules
 
-def find_in_bag(bag, color, rules):
-    if color in rules[bag]: return True
-    if not rules[bag]: return False
 
-    return any([find_in_bag(bag, color, rules) for bag in rules[bag]])
+def find_in_bags(color, rules):
+    count = 0
 
+    for bag in rules:
+        discovered = []
+        discovered.extend(rules[bag])
 
+        while discovered:
+            x, _ = discovered.pop(0)
+            if x == 'shiny gold':
+                count += 1
+                break
+            discovered.extend(rules[x])
 
-def create_bag_tree(bag, rules):
-    for rule in rules[bag.color]:
-        for _ in range(rules[bag.color][rule]):
-            child = create_bag_tree(Bag(rule), rules)
-            bag.append(child)
+    return count
 
-    return bag
+def count_bags(bag, rules):
+    discovered = [bag]
+    total = 0
+
+    while discovered:
+        bag, amount = discovered.pop(0)
+        for x, x_amt in rules[bag]:
+            discovered.append((x, amount * x_amt))
+        total += amount
+
+    return total - 1
 
 
 def solve_a(rules):
-    bags = ([bag for bag in rules if find_in_bag(bag, "shiny gold", rules)])
-    return len(bags)
+    return find_in_bags("shiny gold", rules)
 
 def solve_b(rules):
-    bag = create_bag_tree(Bag("shiny gold"), rules)
-    return bag.count()
+    return count_bags(("shiny gold", 1), rules)
 
 
 data = [line.strip() for line in open('input')]