AoC day 19

code.py

from functools import partial
from itertools import count
import re

INPUT = """Al => ThF
Al => ThRnFAr
B => BCa
B => TiB
B => TiRnFAr
Ca => CaCa
Ca => PB
Ca => PRnFAr
Ca => SiRnFYFAr
Ca => SiRnMgAr
Ca => SiTh
F => CaF
F => PMg
F => SiAl
H => CRnAlAr
H => CRnFYFYFAr
H => CRnFYMgAr
H => CRnMgYFAr
H => HCa
H => NRnFYFAr
H => NRnMgAr
H => NTh
H => OB
H => ORnFAr
Mg => BF
Mg => TiMg
N => CRnFAr
N => HSi
O => CRnFYFAr
O => CRnMgAr
O => HP
O => NRnFAr
O => OTi
P => CaP
P => PTi
P => SiRnFAr
Si => CaSi
Th => ThCa
Ti => BP
Ti => TiTi
e => HF
e => NAl
e => OMg"""

REPLACEMENTS = {re.match(r'(\w+) => (\w+)', line).groups() for line in INPUT.splitlines()}
RUDOLPH_MOLECULE = "CRnSiRnCaPTiMgYCaPTiRnFArSiThFArCaSiThSiThPBCaCaSiRnSiRnTiTiMgArPBCaPMgYPTiRnFArFArCaSiRnBPMgArPRnCaPTiRnFArCaSiThCaCaFArPBCaCaPTiTiRnFArCaSiRnSiAlYSiThRnFArArCaSiRnBFArCaCaSiRnSiThCaCaCaFYCaPTiBCaSiThCaSiThPMgArSiRnCaPBFYCaCaFArCaCaCaCaSiThCaSiRnPRnFArPBSiThPRnFArSiRnMgArCaFYFArCaSiRnSiAlArTiTiTiTiTiTiTiRnPMgArPTiTiTiBSiRnSiAlArTiTiRnPMgArCaFYBPBPTiRnSiRnMgArSiThCaFArCaSiThFArPRnFArCaSiRnTiBSiThSiRnSiAlYCaFArPRnFArSiThCaFArCaCaSiThCaCaCaSiRnPRnCaFArFYPMgArCaPBCaPBSiRnFYPBCaFArCaSiAl"


def expand(molecule, start=""):
    molecules = set()
    for match, replacement in REPLACEMENTS:
        if molecule[:len(match)] == match:
            molecules.add(start + replacement + molecule[len(match):])

    if len(molecule) > 0:
        molecules |= expand(molecule[1:], start + molecule[0])

    return molecules


print "Part One:", len(expand(RUDOLPH_MOLECULE))

REPLACEMENTS = {(r2, r1) for (r1, r2) in REPLACEMENTS}

expansions = {RUDOLPH_MOLECULE: 0}

for i in count():
    next_best = sorted(expansions.keys(), key=len)[0]
    generation = expansions[next_best] + 1
    for expansion in expand(next_best):
        expansions[expansion] = generation
    if 'e' in expansions:
        print "Part Two:", expansions['e']
        break
    del expansions[next_best]