It was hard to remember trees :c but it worked

2020/day_7/handy_haversacks.py

@@ -0,0 +1,68 @@
+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")]
+
+        color = color.replace('bags', '').strip()
+        rule = rule.replace('.', '').split(', ')
+
+        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
+
+    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 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 bag
+
+
+def solve_a(rules):
+    bags = ([bag for bag in rules if find_in_bag(bag, "shiny gold", rules)])
+    return len(bags)
+
+def solve_b(rules):
+    bag = create_bag_tree(Bag("shiny gold"), rules)
+    return bag.count()
+
+
+data = [line.strip() for line in open('input')]
+rules = process_rules(data)
+print(solve_a(rules))
+print(solve_b(rules))
+