advent_14.py

from collections import defaultdict

memo = {}

def key(char, next_char, steps):
  return '{0}{1}{2}'.format(char, next_char, steps)

def combine(current, left, right):
  count = defaultdict(lambda: 0, left)
  count[current] += 1
  for key in right:
    count[key] += right[key]
  return count

def evaluate(char, next_char, rules, steps):
  if steps == 0:
    return defaultdict(lambda: 0)
  current_key = key(char, next_char, steps)
  if current_key in memo:
    return memo[current_key]

  result = rules[char][next_char]
  before = evaluate(char, result, rules, steps - 1)
  after = evaluate(result, next_char, rules, steps - 1)
  counts = combine(result, before, after)
  memo[current_key] = counts
  return counts

if __name__ == '__main__':
  rules = defaultdict(dict)
  with open('input.txt') as data:
    template = data.readline().strip()
    for line in data:
      line = line.strip()
      if not line:
        continue
      source, target = line.split(' -> ')
      rules[source[0]][source[1]] = target
    frequency = defaultdict(lambda: 0)
    totals = defaultdict(lambda: 0)
    for i in range(0, len(template) - 1):
      char = template[i]
      next_char = template[i + 1]
      count = evaluate(char, next_char, rules, 40)
      totals = combine(char, totals, count)
    totals[template[-1]] += 1
    sorted_elements = sorted(totals.keys(), key = lambda item: totals[item])
    print(totals[sorted_elements[-1]] - totals[sorted_elements[0]])