import re from collections import defaultdict def parse_atoms(molecule): return re.findall(r"[A-Z][a-z]*", molecule) decay = defaultdict(list) with open("input") as f: line = f.readline() while line != "\n": atom, decays_to = line.strip().split(" => ") decay[atom].append(decays_to) line = f.readline() molecule = f.readline() print(decay) print() print(molecule) atoms = parse_atoms(molecule) next_molecules = set() for idx, atom in enumerate(atoms): prefix = "".join(atoms[:idx]) suffix = "".join(atoms[idx + 1 :]) for decay_molecule in decay[atom]: descendant = prefix + decay_molecule + suffix next_molecules.add(descendant) print(len(next_molecules))