day07/bags.py

   1 #!/usr/bin/python3
   2
   3 import re
   4
   5 contained_by={}
   6 contains={}
   7
   8 for rule in open("input.txt", "r"):
   9     rule=rule.rstrip()
  10     print("Rule: {}".format(rule))
  11     m=re.match("^(.*?) bag(s|) contain (.*)$", rule)
  12     if m:
  13         bag=m.group(1)
  14         contain=m.group(3).split(", ")
  15         for contained in contain:
  16             m2=re.match('^([0-9]*) (.*?) bag(s|).*$', contained)
  17             if m2 is not None:
  18                 bag_count=int(m2.group(1))
  19                 bag_name=m2.group(2)
  20                 if not bag in contains:
  21                     contains[bag]=[]
  22                 contains[bag].append((bag_count, bag_name))
  23                 if not bag_name in contained_by:
  24                     contained_by[bag_name] = []
  25                 if not bag in contained_by[bag_name]:
  26                     contained_by[bag_name].append(bag)
  27
  28 looking_for="shiny gold"
  29
  30 bags=[]
  31 for bag in contained_by[looking_for]:
  32     if bag not in bags:
  33         bags.append(bag)
  34
  35 print(bags)
  36
  37 new_bags=True
  38 new_bag_list=[b for b in bags]
  39 new_new_bag_list=[]
  40 while new_bags:
  41     for bag1 in new_bag_list:
  42         if bag1 not in contained_by:
  43             continue
  44         for bag2 in contained_by[bag1]:
  45             if bag2 not in bags:
  46                 new_new_bag_list.append(bag2)
  47                 bags.append(bag2)
  48     if (len(new_new_bag_list) > 0):
  49         new_bag_list=[b for b in new_new_bag_list]
  50         new_new_bag_list=[]
  51         new_bags=True
  52     else:
  53         new_bag_list=[]
  54         new_bags=False
  55
  56 print("Found {} bags that can contain {} bags".format(len(bags), looking_for))
  57
  58 def get_count(bag_name):
  59     count=0
  60     if not bag_name in contains:
  61         return 0
  62     for (num,bag) in contains[bag_name]:
  63         count+=(get_count(bag)+1)*num
  64     return count
  65
  66 # now looking for how many bags in bags we're looking for
  67 print("Contains {} bags".format(get_count(looking_for)))