aoc-2020/21.py

from itertools import product
lines = []

with open("21.input") as f:
    for line in f.readlines():
        lines.append(line.strip())

pairs = []
all_allergens = set()
all_ingredients = set()

for line in lines:
    [ingredients, allergens] = line.split(" (contains ")
    allergens = allergens[:-1] # Remove last parentheses
    allergens = allergens.split(", ")
    ingredients = ingredients.split(" ")
    all_allergens = all_allergens.union(set(allergens))
    all_ingredients = all_ingredients.union(set(ingredients))
    pairs.append((ingredients, allergens))


# Returns True if a implies b
def impl(a, b):
    if not a:
        return True
    else:
        return b

possible_combinations = [(a, i) for (a, i) in product(all_allergens, all_ingredients) if all([impl(a in als, i in ins) for (ins, als) in pairs])]

possible_options = {}
for (a, i) in possible_combinations:
    l = possible_options.get(a, [])
    l.append(i)
    possible_options[a] = l

unsafe_ingredients = [i for (a, i) in possible_combinations]
safe_ingredients = [i for i in all_ingredients if i not in unsafe_ingredients]

print("Answer 1:", len([i for (ins, als) in pairs for i in ins if i not in unsafe_ingredients]))

danger_list = []

while len(possible_options) > 0:
    for (k, v) in possible_options.items():
        if len(v) == 1:
            ingredient = v[0]
            danger_list.append((k, ingredient))
            for (k2, v2) in possible_options.items():
                if ingredient in v2:
                    v2.remove(ingredient)
            del possible_options[k]
            break

danger_list.sort()

print("Answer 2:", ",".join([b for (a, b) in danger_list]))