AOC2021 Day 14 (Python)

day14_tests.py

import re


def read_file(file):
    i = iter(file)
    start_string = next(i).strip()

    next(i)
    inserts = dict()

    while True:
        try:
            (a, b), insert = re.search(r'([A-Z][A-Z]) -> ([A-Z])', next(i)).groups()
        except StopIteration:
            break
        inserts[(a, b)] = insert

    return start_string, inserts


def pairs(sequence):
    return ((sequence[i], sequence[i + 1]) for i in range(len(sequence) - 1))


def initial_counts(items):
    counts = dict()
    for item in items:
        counts[item] = counts.get(item, 0) + 1
    return counts


def update_counts(inserts, pair_counts, element_counts):
    next_pair_counts = dict()
    next_element_counts = dict(element_counts)

    for (a, b), count in pair_counts.items():
        insert = inserts.get((a, b), None)
        if insert:
            next_pair_counts[(a, insert)] = next_pair_counts.get((a, insert), 0) + count
            next_pair_counts[(insert, b)] = next_pair_counts.get((insert, b), 0) + count
            next_element_counts[insert] = next_element_counts.get(insert, 0) + count
        else:
            next_pair_counts[(a, b)] = next_pair_counts.get((a, b), 0) + count

    return next_pair_counts, next_element_counts


def perform_inserts(start_string, inserts, times):
    element_counts = initial_counts(start_string)
    pair_counts = initial_counts(pairs(start_string))

    for _ in range(times):
        pair_counts, element_counts = update_counts(inserts, pair_counts, element_counts)

    most_common = max(element_counts.values())
    least_common = min(element_counts.values())
    return most_common - least_common


def test_inserts():
    test_input = """
NNCB

CH -> B
HH -> N
CB -> H
NH -> C
HB -> C
HC -> B
HN -> C
NN -> C
BH -> H
NC -> B
NB -> B
BN -> B
BB -> N
BC -> B
CC -> N
CN -> C
    """.strip().split('\n')

    start_string, inserts = read_file(test_input)
    assert perform_inserts(start_string, inserts, 10) == 1588

    with open("puzzle_inputs/day14.txt") as file:
        start_string, inserts = read_file(file)

    print()
    print(perform_inserts(start_string, inserts, 10))
    print(perform_inserts(start_string, inserts, 40))