65 lines
1.5 KiB
Python
65 lines
1.5 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import functools
|
|
import sys
|
|
|
|
input_file = sys.argv[1]
|
|
|
|
with open(input_file) as fd:
|
|
lines = [line.rstrip() for line in fd.readlines()]
|
|
|
|
numeric_keypad = ["789", "456", "123", " 0A"]
|
|
directional_keypad = [" ^A", "<v>"]
|
|
|
|
vec = tuple[int, int]
|
|
|
|
complexity = int(sys.argv[2])
|
|
keypads = [numeric_keypad] + ([directional_keypad] * complexity)
|
|
|
|
|
|
@functools.cache
|
|
def press(buts: str, depth: int) -> int:
|
|
if depth == len(keypads):
|
|
return len(buts)
|
|
|
|
keypad = keypads[depth]
|
|
but_poss: dict[str, vec] = dict()
|
|
for i, line in enumerate(keypad):
|
|
for j, but in enumerate(line):
|
|
but_poss[but] = i, j
|
|
i, j = but_poss["A"]
|
|
ai, aj = but_poss[" "]
|
|
|
|
depth += 1
|
|
nums = 0
|
|
for but in buts:
|
|
nnbuts = ""
|
|
ni, nj = but_poss[but]
|
|
crossing_gap = ((i == ai) and (nj == aj)) or ((ni == ai) and (j == aj))
|
|
if nj < j:
|
|
nnbuts += "<" * (j - nj)
|
|
if ni < i:
|
|
nnbuts += "^" * (i - ni)
|
|
elif ni > i:
|
|
nnbuts += "v" * (ni - i)
|
|
if nj > j:
|
|
nnbuts += ">" * (nj - j)
|
|
i, j = ni, nj
|
|
if crossing_gap:
|
|
nnbuts = nnbuts[::-1]
|
|
nnbuts += "A"
|
|
nums += press(nnbuts, depth)
|
|
return nums
|
|
|
|
|
|
score = 0
|
|
for code in lines:
|
|
print("Code", code)
|
|
topresses = press(code, 0)
|
|
|
|
numpart = int("0" + code.replace("A", ""))
|
|
print(f"{topresses=} * {numpart=}")
|
|
score += topresses * numpart
|
|
|
|
print(score)
|