AoC 2021 Day 8

import os
from collections import Counter
import timeit


def read_and_init(fn):
    with open(fn, 'r') as f:
        lines = f.read().splitlines()
    entries = {}
    i = 0
    for line in lines:
        patterns, digits = line.split(' | ')
        entries[i] = {'patterns': patterns.split(), 'digits': digits.split()}
        i += 1
    return entries


def part1(entries: list) -> int:
    segments = {0: 6, 1: 2, 2: 5, 3: 5, 4: 4, 5: 5, 6: 6, 7: 3, 8: 7, 9: 6}
    counter = Counter(segments.values())
    uniques = [key for key in counter if counter[key] == 1]
    total_uniques = 0
    for entry in entries:
        digit_lengths = [len(d) for d in entries[entry]['digits']]
        entry_uniques = [dl for dl in digit_lengths if dl in uniques]
        total_uniques += len(entry_uniques)
    return total_uniques


def part2(entries):
    readouts = []
    numerals = {i: '' for i in range(10)}
    for entry in entries:
        patterns = entries[entry]['patterns']
        digits = [''.join(sorted(digit)) for digit in entries[entry]['digits']]
        patterns.sort()
        # start with uniques
        numerals[1] = [pattern for pattern in patterns if len(pattern) == 2][0]
        numerals[4] = [pattern for pattern in patterns if len(pattern) == 4][0]
        numerals[7] = [pattern for pattern in patterns if len(pattern) == 3][0]
        numerals[8] = [pattern for pattern in patterns if len(pattern) == 7][0]
        fives = [pattern for pattern in patterns if len(pattern) == 5]
        sixes = [pattern for pattern in patterns if len(pattern) == 6]
        numerals[6] = [pattern for pattern in sixes if len(
            set(numerals[1]) - (set(pattern) & set(numerals[1]))) == 1][0]
        sixes.remove(numerals[6])
        numerals[5] = [pattern for pattern in fives if len(set(numerals[6]) - set(pattern)) == 1][0]
        fives.remove(numerals[5])
        numerals[2] = [pattern for pattern in fives if len(set(pattern) - set(numerals[1])) == 4][0]
        fives.remove(numerals[2])
        numerals[3] = fives.pop()
        numerals[0] = [pattern for pattern in sixes if len(set(pattern) - set(numerals[4])) == 3][0]
        sixes.remove(numerals[0])
        numerals[9] = sixes.pop()
        swapped_numerals = {''.join(sorted(v)): k for k, v in numerals.items()}
        decoded_digits = [swapped_numerals[digit] for digit in digits]
        readout = int(''.join(map(str, decoded_digits)))
        readouts.append(readout)
    return sum(readouts)


es = read_and_init(os.path.dirname(__file__) + '\\in.txt')
result = part1(es)
print('Part 1:', result)
result = part2(es)
print('Part 2:', result)