aoc-2020/19.py

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))