98 lines
2.7 KiB
Python
98 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
|
|
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]
|
|
directions = {
|
|
"^": (-1, 0), # ^ North
|
|
">": (0, 1), # > East
|
|
"v": (1, 0), # v South
|
|
"<": (0, -1), # < West
|
|
}
|
|
directional_keypad_buttons = tuple(directions.keys()) + ("A",)
|
|
|
|
complexity = int(sys.argv[2])
|
|
keypads = [directional_keypad] * complexity + [numeric_keypad]
|
|
|
|
|
|
def but_pos(but: str, keypad: list[str]) -> vec:
|
|
for i, line in enumerate(keypad):
|
|
if but in line:
|
|
return i, line.index(but)
|
|
raise IndexError("No such button")
|
|
|
|
|
|
def in_bounds(i: int, j: int, keypad: list[str]) -> bool:
|
|
if j not in range(3) or i not in range(len(keypad)):
|
|
return False
|
|
return keypad[i][j] != " "
|
|
|
|
last_but = "A"
|
|
all_a = [but_pos("A", keypad) for keypad in keypads]
|
|
|
|
score = 0
|
|
for code in lines:
|
|
# print("Code", code)
|
|
topresses = 0
|
|
for desir_but in code:
|
|
# print("Button", desir_but)
|
|
all_a[-1] = but_pos(last_but, keypads[-1])
|
|
start_poss = tuple(all_a)
|
|
all_a[-1] = but_pos(desir_but, keypads[-1])
|
|
if len(keypads) > 1:
|
|
all_a[-2] = but_pos("^", keypads[-2])
|
|
desir_poss = tuple(all_a)
|
|
|
|
stack = {start_poss}
|
|
seen = set()
|
|
presses = 0
|
|
while desir_poss not in stack:
|
|
# print("Press", presses, stack)
|
|
presses += 1
|
|
nstack = set()
|
|
for poss in stack:
|
|
for but in directional_keypad_buttons:
|
|
# Find which keypad this will move
|
|
k = 0
|
|
while but == "A" and k < len(keypads) - 1:
|
|
i, j = poss[k]
|
|
but = keypads[k][i][j]
|
|
k += 1
|
|
|
|
# Do not press the final keypad
|
|
if k == len(keypads) - 1 and but == "A":
|
|
continue
|
|
|
|
# Move
|
|
direction = directions[but]
|
|
i, j = poss[k]
|
|
ii, jj = i + direction[0], j + direction[1]
|
|
if not in_bounds(ii, jj, keypads[k]):
|
|
continue
|
|
|
|
# Ensure we haven't been in this state before
|
|
state = poss[:k] + ((ii, jj),) + poss[k + 1 :]
|
|
if state in seen:
|
|
continue
|
|
seen.add(state)
|
|
|
|
# print(" Kept", state)
|
|
nstack.add(state)
|
|
|
|
stack = nstack
|
|
topresses += presses + 0
|
|
last_but = desir_but
|
|
|
|
score += topresses
|
|
|
|
print(score)
|
|
|