Advent of Code 2021 day 14 part 2

AoC_2021_14_2.py

import re
from collections import Counter, defaultdict

with open("input_14.txt") as f:
    template = f.readline().strip()
    rules = {
        tuple(pair): insert
        for pair, insert in re.findall(r"(..) -> (.)", f.read())
    }
elem_frst = template[0]
elem_last = template[-1]

pair_cnt = Counter(zip(template, template[1:]))
for _ in range(40):
    tmp = defaultdict(int)
    for (ch0, ch1), ins in rules.items():
        tmp[ch0, ins] += pair_cnt[ch0, ch1]
        tmp[ins, ch1] += pair_cnt[ch0, ch1]
    pair_cnt = tmp

elem_cnt = defaultdict(float)
for (ch0, ch1), cnt in pair_cnt.items():
    elem_cnt[ch0] += cnt / 2
    elem_cnt[ch1] += cnt / 2
elem_cnt[elem_frst] += 0.5
elem_cnt[elem_last] += 0.5

print(int(max(elem_cnt.values()) - min(elem_cnt.values())))