from itertools import product lines = [] with open("19.input") as f: for line in f.readlines(): lines.append(line.strip()) rules = {} i = 0 while lines[i] != "": [id, rule] = lines[i].split(":") if '"' in rule: rules[int(id)] = rule.strip().replace('"', "") i += 1 continue alternatives = rule.strip().split("|") for j in range(0, len(alternatives)): alternatives[j] = [int(n) for n in alternatives[j].strip().split(" ")] rules[int(id)] = alternatives i += 1 words = lines[i+1:] def matching_words(rule_id): rule = rules[rule_id] if type(rule) is str: return set(rule) words = set() for alt in rule: if len(alt) == 1: words = words.union(matching_words(alt[0])) elif len(alt) == 2: for (x, y) in product(matching_words(alt[0]), matching_words(alt[1])): words.add(x + y) else: print("Invalid length: ", len(alt)) return words allowed_words = matching_words(0) matched_words = [word for word in words if word in allowed_words] print("Answer 1:", len(matched_words)) special_42 = matching_words(42) special_31 = matching_words(31) def try_match(word): count_42 = 0 count_31 = 0 original_word = word while True: for special in special_42: if word.find(special) == 0: word = word[len(special):] count_42 += 1 break else: break while True: for special in special_31: if word.find(special) == 0: word = word[len(special):] count_31 += 1 break else: break # Must match the entire word if len(word) != 0: return False return count_31 >= 1 and count_42 >= count_31 + 1 matched_words = [word for word in words if try_match(word)] print("Answer 2:", len(matched_words))