Advent of Code Day 21

"""Day 22 of Advent of Code 2020 Solution"""
from collections import Counter


def allergen_io(file_location):
    word_count = Counter()
    word_sets = {}
    with open(file_location, 'r') as f:
        for line in f.readlines():
            line = line.strip()[:-1].split('(')
            line[0] = line[0].strip().split(' ')
            line[1] = line[1][9:].strip().split(', ')
            word_count.update(line[0])
            for allergen in line[1]:
                if allergen in word_sets:
                    word_sets[allergen].intersection_update(line[0])
                else:
                    word_sets[allergen] = {*line[0]}
    return word_count, word_sets


def part_a(file_location):
    word_count, word_sets = allergen_io(file_location)
    allergen_set = set()
    for allergen in word_sets.values():
        allergen_set.update(allergen)
    n_non_allergen_words = 0
    for word in word_count:
        if word not in allergen_set:
            n_non_allergen_words += word_count[word]
    return n_non_allergen_words


def part_b(file_location):
    _, word_sets = allergen_io(file_location)
    while not all([len(v) == 1 for v in word_sets.values()]):
        allergens = sorted([*word_sets.keys()], key=lambda x: len(word_sets[x]))
        for i, allergen_1 in enumerate(allergens):
            for allergen_2 in allergens[i + 1:]:
                if len(word_sets[allergen_1]) == 1:
                    word_sets[allergen_2].difference_update(word_sets[allergen_1])
    return ','.join([word_sets[key].pop() for key in sorted([*word_sets.keys()])])


if __name__ == '__main__':
    file_location = r"data\day21.txt"
    print(part_a(file_location))
    print(part_b(file_location))