AoC 21-14

from collections import Counter

file = "inputs/2021/14/data.txt"
with open(file, "r") as f:
    template = f.readline().strip()
    f.readline()
    rules = f.read().split("\n")
    rules = {rule.split(" -> ")[0]: rule.split(" -> ")[1] for rule in rules}

def polymerize(n_polymerizations, template, rules):

    # Initialize pair_counts
    pair_counts = Counter()
    for idx in range(len(template) - 1):
        pair_counts[template[idx:idx + 2]] += 1

    # Run polymerization
    for polymerization_idx in range(n_polymerizations):
        new_pair_counts = Counter()
        for pair, count in pair_counts.items():
            insertion = rules[pair]
            new_pair_counts[pair[0] + insertion] += count
            new_pair_counts[insertion + pair[1]] += count

        pair_counts = new_pair_counts

    # Calculate result
    letter_counts = Counter()
    for pair, count in pair_counts.items():
        letter_counts[pair[0]] += count
        letter_counts[pair[1]] += count

    # Adjust letter counts for duplication
    for letter in letter_counts:
        letter_counts[letter] = letter_counts[letter] - letter_counts[letter] // 2

    return letter_counts


part_1_base_counts = polymerize(10, template, rules)
part_1 = max(part_1_base_counts.values()) - min(part_1_base_counts.values())
print(f"Part 1: {part_1:,d}")

part_2_base_counts = polymerize(40, template, rules)
part_2 = max(part_2_base_counts.values()) - min(part_2_base_counts.values())
print(f"Part 2: {part_2:,d}")