aoc2021-14.py

import re
import sys
from collections import defaultdict


def parse_data():
    pairs = defaultdict(int)
    data = sys.stdin
    polymer = data.readline().strip()
    rules = {}

    for line in data:
        line = line.strip()
        if not line:
            continue
        match = re.match(r"(\w+) -> (\w+)", line)
        assert match
        pair, e = match.groups()
        rules[pair] = e

    for idx in range(len(polymer)):
        # When idx reaches len(polymer) - 1, we'll get a single char substring
        # here.  That's expected and needed to ensure counts at the end are
        # correct.
        pair = polymer[idx : idx + 2]
        pairs[pair] += 1

    return pairs, rules


def replace(pairs, rules):
    result = defaultdict(int)
    for pair, count in pairs.items():
        rule = rules.get(pair)
        if rule:
            pair_1 = pair[0] + rule
            pair_2 = rule + pair[1]
            result[pair_1] += count
            result[pair_2] += count
        else:
            result[pair] += count
    return result


def get_solution(pairs):
    counts = defaultdict(int)
    for pair, count in pairs.items():
        counts[pair[0]] += count

    return max(counts.values()) - min(counts.values())


pairs, rules = parse_data()


for _ in range(10):
    pairs = replace(pairs, rules)

print(f"Part 1: {get_solution(pairs)}")

for _ in range(30):
    pairs = replace(pairs, rules)

print(f"Part 2: {get_solution(pairs)}")