21_keypad.py 1.1 KB

123456789101112131415161718192021222324252627282930313233343536373839
  1. #!/usr/bin/env python3
  2. import sys
  3. from functools import cache
  4. from itertools import pairwise
  5. NUMERIC = '789' '456' '123' ' 0A'
  6. DIRECTIONAL = ' ^A' '<v>'
  7. def walk(keypad, x, y, path):
  8. i = y * 3 + x
  9. for direction in path:
  10. i += (-1, 1, -3, 3)['<>^v'.index(direction)]
  11. yield keypad[i]
  12. def paths_between(keypad, start, end):
  13. y1, x1 = divmod(keypad.index(start), 3)
  14. y2, x2 = divmod(keypad.index(end), 3)
  15. hor = '<>'[x2 > x1] * abs(x2 - x1)
  16. ver = '^v'[y2 > y1] * abs(y2 - y1)
  17. for path in {hor + ver, ver + hor}:
  18. if ' ' not in walk(keypad, x1, y1, path):
  19. yield path + 'A'
  20. @cache
  21. def cost_between(keypad, start, end, links):
  22. return min(cost(DIRECTIONAL, path, links - 1)
  23. for path in paths_between(keypad, start, end)) if links else 1
  24. @cache
  25. def cost(keypad, keys, links):
  26. return sum(cost_between(keypad, a, b, links)
  27. for a, b in pairwise('A' + keys))
  28. def complexity(code, robots):
  29. return cost(NUMERIC, code, robots + 1) * int(code[:-1])
  30. codes = sys.stdin.read().split()
  31. print(sum(complexity(code, 2) for code in codes))
  32. print(sum(complexity(code, 25) for code in codes))