AoC day 14: polymer, another exponential growth one

day14.py

#!/usr/bin/env python3
#
# https://adventofcode.com/2021/day/14

import sys, os,math
from collections import deque, defaultdict

def to_pairs(template):
    result = defaultdict(int)
    for i in range(len(template)-1):
        result["".join(template[i:i+2])] += 1
    return result

def to_singles(template):
    result = defaultdict(int)
    for i,ch in enumerate(template):
        result[ch] += 1
    return result

def step(pairs, singles, rules):
    scancopy = pairs.copy()
    for rule, addition in rules.items():
        found_qty = scancopy[rule]
        if found_qty:
            #print("found",found_qty,"of",rule)
            left=rule[0]+addition
            right=addition+rule[1]
            #print("replacing with equal amounts",left,right), 
            pairs[rule] -= found_qty
            pairs[left]  += found_qty
            pairs[right] += found_qty
            #print("adding",found_qty,"of",addition)
            singles[addition] += found_qty

def parse_lines(lines):
    rules = {}
    for line in lines:
        if "->" in line:
            k, v = line.rstrip().split(" -> ")
            rules[k]=v
        elif line.rstrip():
            template = (line.rstrip())

    return (rules, template)

def part1(lines):
    rules, template = parse_lines(lines)
    print("rules=%r\ntemplate=%r" % (rules, template))

    current = to_pairs(template)
    singles = to_singles(template)
    assert singles is not None
    print(current)
    for i in range(10):
        step(current, singles, rules)
        print(".",end='', flush=True)
    print("part1:", score(singles))
    for i in range(30):
        step(current, singles, rules)
        print(".",end='',flush=True)
    print("part2:", score(singles))

    
def score(single_letters):
    elements = list(single_letters.keys())
    elements.sort(key=single_letters.__getitem__)
    return single_letters[ elements[-1] ] - single_letters[elements[0]]


if __name__ == '__main__':
    if len(sys.argv)<2:
        print("Usage",sys.argv[0],"filename")
        sys.exit(1)
    with open(sys.argv[1]) as f:
        lines = f.readlines()

    part1(lines)

test.py

#!/usr/bin/env python3

import unittest

from day14 import *

class FooTest(unittest.TestCase):
    def test_to_pairs(self):
        pairs = to_pairs("NNCB")
        self.assertEqual(pairs, {"NN":1, "NC":1, "CB":1})

    def test_to_singles(self):
        singles = to_singles("NNCB")
        self.assertEqual(singles, {"N":2, "C":1, "B":1})

    def test_example_first_step(self):
        template='NNCB'
        rules={'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'}
        singles = to_singles(template)
        pairs = to_pairs(template)
        step(pairs,singles, rules)
        self.assertEqual(dict(singles), {"N":2, "C":2, "B":2, "H":1})
        self.assertEqual(dict(pairs), {"NN":0, "NC":1, "CN":1, "NB":1, "BC":1, "CH":1, "HB":1, "CB":0})


if __name__ == '__main__':
    unittest.main()